Sgp/ex5da

examples/da_cycler/5-ETKF-NeuralGCM.ipynb

@@ -0,0 +1,254 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "d6ea0e40-f74b-4800-b678-4785912e74c9",
+   "metadata": {},
+   "source": [
+    "# 5. Example: DataAssimBench ensemble data assimilation using the NeuralGCM\n",
+    "\n",
+    "The NeuralGCM takes input from an ERA5-like format and returns a forecast at roughly 1-hour intervals. Here we use an initial ensemble drawn from ERA5 and assimilate a single observation sampled from the ERA5 'truth'."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3103069a-265a-4887-9ae5-45b3bef9c0a3",
+   "metadata": {},
+   "source": [
+    "#### Get required packages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4329348a-a718-4f85-ad26-27037acb3f79",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%capture\n",
+    "\n",
+    "import sys\n",
+    "\n",
+    "# Make sure that the gcsfs package is installed in the current Jupyter kernel\n",
+    "!{sys.executable} -m pip install gcsfs\n",
+    "\n",
+    "# Make sure the dinosaur dynamical core is installed for NeuralGCM\n",
+    "!{sys.executable} -m pip install dinosaur"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ba1f214f-ef40-4f46-90f5-b51bbf17c85c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import dabench as dab\n",
+    "from datetime import datetime\n",
+    "import xarray as xr\n",
+    "import os\n",
+    "import matplotlib.pyplot as plt"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6193c156-b946-49c0-8ae2-93c4124dce99",
+   "metadata": {},
+   "source": [
+    "#### Generate an initial ensemble"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "598e83d0-fb65-4547-8278-b524ae58758c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Point to the location where the ERA5 dataset is stored on GCP\n",
+    "gcp_era5_path = \"gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3\"\n",
+    "\n",
+    "# Specify the desired lagged ensemble sampling strategy\n",
+    "sample_strategy = 'multi_year'\n",
+    "\n",
+    "# Specify the target date to which the ensemble will be recentered\n",
+    "target_date = \"20240921Z0\"\n",
+    "\n",
+    "# Specify the desired ensemble size\n",
+    "ensemble_size = 3\n",
+    "\n",
+    "# Store locally or on cloud via zarr\n",
+    "output_destination = f'./zarr_{target_date}_{ensemble_size}mem'\n",
+    "\n",
+    "# Only generate if you have not yet processed this date\n",
+    "if not os.path.exists(output_destination):\n",
+    "    # Set up the input arguments\n",
+    "    input_args = {\n",
+    "        'date_format': \"%Y%m%dZ%H\",\n",
+    "        'atmosphere_ensemble_s3_key': output_destination,\n",
+    "        'target_date': datetime.strptime(target_date,\"%Y%m%dZ%H\"),\n",
+    "        'sample_strategy': sample_strategy,\n",
+    "        'start_date': datetime.strptime(target_date,\"%Y%m%dZ%H\"),\n",
+    "        'ensemble_size': ensemble_size,\n",
+    "        'era5_path': gcp_era5_path,\n",
+    "    }\n",
+    "    \n",
+    "    # Generate the initial ensemble, and store to a zarr file\n",
+    "    dab.dasupport.GenEra5Ens(**input_args)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fabeb5cd-0bca-4112-bd30-cc7540019f49",
+   "metadata": {},
+   "source": [
+    "#### Check the initial ensemble"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cfdfe669-1071-4124-8a7f-0f430bc41d7e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Open the ensemble dataset stored in zarr format as an xarray dataset\n",
+    "ds = xr.open_dataset(output_destination, engine='zarr')\n",
+    "ds"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7ac987f5-a55a-46ca-8b34-7d7a1f8a0c40",
+   "metadata": {},
+   "source": [
+    "#### Test run the NeuralGCM model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "26e7f4fc-242d-479c-a759-eefa4eea591f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Store the ensemble mean to use as input to the NeuralGCM\n",
+    "output_destination_mean = f'{output_destination}_mean'\n",
+    "ds.mean(dim='member').to_zarr(output_destination_mean, mode='w')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0685c700-c4b8-4e59-a059-f88ee37cb5fd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_input_args = {\n",
+    "    'era5_path': output_destination_mean,\n",
+    "    'atm_res': '2_8',\n",
+    "    'forecast_hours': 1,\n",
+    "    'data_stride': 1,\n",
+    "    'start_time': datetime.strptime(target_date,\"%Y%m%dZ%H\"),\n",
+    "    'inner_steps': 1,\n",
+    "    'outer_steps': 1,\n",
+    "}\n",
+    "model = NeuralGCM(params=model_input_args)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eeefb74b-c732-4705-8d28-2afaf56ac30e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Run test forecast\n",
+    "model.full_sequence()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "714a61c4-804d-4c71-aeca-647d1d4f28f4",
+   "metadata": {},
+   "source": [
+    "## Apply data assimilation"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2b694e61-1671-496a-bb0b-a652bbfbfb54",
+   "metadata": {},
+   "source": [
+    "#### Set up the ETKF object"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f89a866e-5d12-4f5c-b0a7-a579604e3733",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Pass the NeuralGCM model object into the DA cycler\n",
+    "da_input_args = {\n",
+    "    system_dim: ,\n",
+    "    ensemble_dim: ensemble_size,\n",
+    "    delta_t: 3600,\n",
+    "    model_obj: model,\n",
+    "    H: ,\n",
+    "    h: ,\n",
+    "}\n",
+    "etkf = dab.ETKF(**da_input_args)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f2c8f984-1ad6-4d94-a91d-cca3bee3caf8",
+   "metadata": {},
+   "source": [
+    "#### Perform a DA cycle"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4f1b6219-2c8c-455d-8d83-27933d6be8b7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "etkf_input_args = {\n",
+    "    input_state: ,\n",
+    "    start_time: target_date,\n",
+    "    obs_vector: ,\n",
+    "    n_cycles: 2,\n",
+    "    obs_error_sd: 1.0, #hPa\n",
+    "    analysis_window: 3600,\n",
+    "    return_forecast=True,\n",
+    "}\n",
+    "etkf.cycle(**etkf_input_args)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.11.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}