mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-26 14:42:04 +00:00
feat: WN2 with GPU GA (#4563)
This commit is contained in:
Damodar Panigrahi
@@ -30,6 +30,8 @@
|
|||||||
/notebooks/community/alphagenome/cloudai_alphagenome_finetune.ipynb @dpanigra
|
/notebooks/community/alphagenome/cloudai_alphagenome_finetune.ipynb @dpanigra
|
||||||
/notebooks/community/weathernext/weathernext_2_early_access_program.ipynb @dpanigra
|
/notebooks/community/weathernext/weathernext_2_early_access_program.ipynb @dpanigra
|
||||||
/notebooks/community/weathernext/weathernext_2_ic_early_access_program.ipynb @dpanigra
|
/notebooks/community/weathernext/weathernext_2_ic_early_access_program.ipynb @dpanigra
|
||||||
|
/notebooks/community/weathernext/weathernext_2_dws.ipynb @dpanigra
|
||||||
|
/notebooks/community/weathernext/weathernext_2_ic_pc.ipynb @dpanigra
|
||||||
/notebooks/community/ml_ops/stage2/get_started_with_visionapi_and_automl.ipynb @mansari
|
/notebooks/community/ml_ops/stage2/get_started_with_visionapi_and_automl.ipynb @mansari
|
||||||
/notebooks/community/neo4j/graph_paysim.ipynb @benofben @laeg
|
/notebooks/community/neo4j/graph_paysim.ipynb @benofben @laeg
|
||||||
/notebooks/community/ml_ops/stage1/get_started_with_visionapi_and_vertex_datasets.ipynb @mansari
|
/notebooks/community/ml_ops/stage1/get_started_with_visionapi_and_vertex_datasets.ipynb @mansari
|
||||||
|
|||||||
@@ -4,10 +4,9 @@
|
|||||||
## Overview
|
## Overview
|
||||||
**Disclaimer:**
|
**Disclaimer:**
|
||||||
|
|
||||||
*Experimental*\
|
This product is subject to the General Service Terms section of the [Service Specific Terms](https://cloud.google.com/terms/service-terms#1). GA products follows standard support. For more information, see the [launch stage descriptions](https://cloud.google.com/products#product-launch-stages). <!-- disableFinding(LINE_OVER_80) -->
|
||||||
This product is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the [Service Specific Terms](https://cloud.google.com/terms/service-terms#1). Pre-GA products are available "as is" and might have limited support. For more information, see the [launch stage descriptions](https://cloud.google.com/products#product-launch-stages). <!-- disableFinding(LINE_OVER_80) -->
|
|
||||||
|
|
||||||
Access to the forecasting capabilities requires application and approval. Users must be added to an allowlist to generate forecasts using this service. Review pricing details at [Vertex AI Custom Training pricing,](https://cloud.google.com/vertex-ai/pricing?hl=en&e=48754805#custom-trained-models) [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) before running. <!-- disableFinding(LINE_OVER_80) -->
|
Please contact your Google Cloud sales person for using WeatherNext 2 on Agent Platform. In case you do not have a salesperson you work with, please submit the form using the Request access button in this model card and a Google Cloud representative will contact you. <!-- disableFinding(LINE_OVER_80) -->
|
||||||
|
|
||||||
**Overview**
|
**Overview**
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,577 @@
|
|||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {
|
||||||
|
"id": "RirpY96_3zL0"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Copyright 2026 Google LLC\n",
|
||||||
|
"\n",
|
||||||
|
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
|
||||||
|
"# you may not use this file except in compliance with the License.\n",
|
||||||
|
"# You may obtain a copy of the License at\n",
|
||||||
|
"\n",
|
||||||
|
"# https://www.apache.org/licenses/LICENSE-2.0\n",
|
||||||
|
"\n",
|
||||||
|
"# Unless required by applicable law or agreed to in writing, software\n",
|
||||||
|
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
|
||||||
|
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
|
||||||
|
"# See the License for the specific language governing permissions and\n",
|
||||||
|
"# limitations under the License."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {
|
||||||
|
"id": "52d0938f"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"# WeatherNext 2 (Using DWS)\n",
|
||||||
|
"<table align=\"left\">\n",
|
||||||
|
" <td style=\"text-align: center\">\n",
|
||||||
|
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/weathernext/weathernext_2_dws.ipynb\">\n",
|
||||||
|
" <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
|
||||||
|
" </a>\n",
|
||||||
|
" </td>\n",
|
||||||
|
" <td style=\"text-align: center\">\n",
|
||||||
|
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%weathernext%2Fweathernext_2_dws.ipynb\">\n",
|
||||||
|
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
|
||||||
|
" </a>\n",
|
||||||
|
" </td>\n",
|
||||||
|
" <td style=\"text-align: center\">\n",
|
||||||
|
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/weathernext/weathernext_2_dws.ipynb\">\n",
|
||||||
|
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
|
||||||
|
" </a>\n",
|
||||||
|
" </td>\n",
|
||||||
|
" <td style=\"text-align: center\">\n",
|
||||||
|
" <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/weathernext/weathernext_2_dws.ipynb\">\n",
|
||||||
|
" <img width=\"32px\"src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
|
||||||
|
" </a>\n",
|
||||||
|
" </td>\n",
|
||||||
|
"</table>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {
|
||||||
|
"id": "DGFzy9MTuYL_"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"## Overview\n",
|
||||||
|
"\n",
|
||||||
|
"This notebook demonstrates running [WeatherNext 2 inference on Google Cloud Vertex AI](https://developers.google.com/weathernext/guides/access-vmg). WeatherNext 2 is Google's latest medium-range probabilistic forecasting model, principally an operational version the FGN model ([published June 2025](https://arxiv.org/abs/2506.10772)). More information is available in the [WeatherNext documentation](https://developers.google.com/weathernext).\n",
|
||||||
|
"\n",
|
||||||
|
"### Objective\n",
|
||||||
|
"\n",
|
||||||
|
"- Configure the model inputs for distributed, multi-host inference on H100 or A100 GPUs.\n",
|
||||||
|
"- Run WeatherNext 2 model forecasts in parallel.\n",
|
||||||
|
"- Visualize forecast results.\n",
|
||||||
|
"\n",
|
||||||
|
"### Costs\n",
|
||||||
|
"\n",
|
||||||
|
"This uses billable components of Google Cloud:\n",
|
||||||
|
"\n",
|
||||||
|
"* Vertex AI\n",
|
||||||
|
"* Cloud Storage\n",
|
||||||
|
"\n",
|
||||||
|
"Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage.\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"## Before you begin\n",
|
||||||
|
"\n",
|
||||||
|
"### Request For GPU Quota\n",
|
||||||
|
"\n",
|
||||||
|
"**WARNING:** Make sure you have sufficient GPU quota allocated for the inference configuration (i.e. `num_samples`) before running Vertex Jobs. Otherwise, some Vertex jobs may run while others will fail which would produce\n",
|
||||||
|
"incomplete results.\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"By default, the quota for GPUs is 0. You can request a higher quota by following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
|
||||||
|
"\n",
|
||||||
|
"You will need to request quota for either **NVIDIA H100 80GB GPUs** or **NVIDIA A100 80GB GPUs** in your selected region. The total number of GPUs you request must be sufficient for your largest planned forecast (i.e., `num_samples`).\n",
|
||||||
|
"\n",
|
||||||
|
"You should request for the following quota:\n",
|
||||||
|
"\n",
|
||||||
|
"- Service: `Vertex AI API`\n",
|
||||||
|
"- Name: `Custom model training preemptible Nvidia A100 80GB GPUs per region` OR `Custom model training preemptible Nvidia H100 GPUs per region`"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {
|
||||||
|
"id": "yt4GxkKcDD7Y"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# @title Install python packages\n",
|
||||||
|
"\n",
|
||||||
|
"# Note that you may need to restart the kernel after this step.\n",
|
||||||
|
"# If so, continue to the next cell after restarting.\n",
|
||||||
|
"\n",
|
||||||
|
"print(\"Installing python packages.\")\n",
|
||||||
|
"\n",
|
||||||
|
"! pip3 install \\\n",
|
||||||
|
" google-cloud-aiplatform==1.129.0 \\\n",
|
||||||
|
" xarray[complete]"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {
|
||||||
|
"id": "dO2mF4CPfpHW1ZKbIsut8r69"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# @title Setup Google Cloud project\n",
|
||||||
|
"\n",
|
||||||
|
"# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
|
||||||
|
"\n",
|
||||||
|
"# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched.\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"BUCKET_URI = \"gs://my-bucket\" # @param {type:\"string\"}\n",
|
||||||
|
"\n",
|
||||||
|
"# @markdown 3. Select a region that has the required GPUs available.\n",
|
||||||
|
"\n",
|
||||||
|
"REGION = \"us-central1\" # @param {type:\"string\"}\n",
|
||||||
|
"\n",
|
||||||
|
"# Import the necessary packages\n",
|
||||||
|
"import datetime\n",
|
||||||
|
"import os\n",
|
||||||
|
"import re\n",
|
||||||
|
"\n",
|
||||||
|
"from google.cloud import aiplatform\n",
|
||||||
|
"\n",
|
||||||
|
"# Get the default cloud project id.\n",
|
||||||
|
"PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
|
||||||
|
"\n",
|
||||||
|
"# Enable the Vertex AI API and Compute Engine API, if not already.\n",
|
||||||
|
"print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
|
||||||
|
"! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
|
||||||
|
"\n",
|
||||||
|
"# Cloud Storage bucket for storing the experiment artifacts.\n",
|
||||||
|
"now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
|
||||||
|
"BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
|
||||||
|
"\n",
|
||||||
|
"if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
|
||||||
|
" raise ValueError(\"GCS Bucket URI is invalid!\")\n",
|
||||||
|
"else:\n",
|
||||||
|
" assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
|
||||||
|
" shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
|
||||||
|
" bucket_region = shell_output[0].strip().lower()\n",
|
||||||
|
" if bucket_region != REGION:\n",
|
||||||
|
" raise ValueError(\n",
|
||||||
|
" f\"Bucket region {bucket_region} is different from notebook region {REGION}\"\n",
|
||||||
|
" )\n",
|
||||||
|
"print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
|
||||||
|
"\n",
|
||||||
|
"STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
|
||||||
|
"# Initialize Vertex AI API.\n",
|
||||||
|
"aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
|
||||||
|
"\n",
|
||||||
|
"# Utility functions\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def get_job_name_with_datetime(prefix: str) -> str:\n",
|
||||||
|
" return prefix + datetime.datetime.now().strftime(\"_%Y%m%d_%H%M%S\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {
|
||||||
|
"id": "22BW43yjps8D"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# @title Configure Model Parameters\n",
|
||||||
|
"# @markdown Configure the hardware and input parameters for the WeatherNext 2 forecast.\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# @markdown ### Hardware Configuration for Distributed Inference\n",
|
||||||
|
"# @markdown - **`machine_type`**: Select a valid machine type. `a3-highgpu` series use NVIDIA H100 80GB GPUs. `a2-ultragpu` series use NVIDIA A100 80GB GPUs.\n",
|
||||||
|
"# @markdown - **`num_samples`**: The total number of ensemble members to generate.\n",
|
||||||
|
"# @markdown The number of machine replicas will be calculated automatically (`num_samples` / GPUs per machine). **Therefore, `num_samples` must be a multiple of the number of GPUs in your selected `machine_type`.**\n",
|
||||||
|
"# @markdown - **`scheduling_strategy`**: The [strategy](https://cloud.google.com/vertex-ai/docs/reference/rest/v1beta1/CustomJobSpec#Strategy) used to acquire machines for the job. Defaults to [Dynamic Workload Scheduler](https://docs.cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws) (FLEX_START).\n",
|
||||||
|
"machine_type = \"a3-highgpu-1g\" # @param [\"a3-highgpu-1g\", \"a3-highgpu-2g\", \"a3-highgpu-4g\", \"a3-highgpu-8g\", \"a2-ultragpu-1g\", \"a2-ultragpu-2g\", \"a2-ultragpu-4g\", \"a2-ultragpu-8g\"]\n",
|
||||||
|
"num_samples = 8 # @param {type:\"integer\"}\n",
|
||||||
|
"scheduling_strategy = \"FLEX_START\" # @param [\"FLEX_START\", \"SPOT\", \"STANDARD\"]\n",
|
||||||
|
"\n",
|
||||||
|
"# @markdown ### Forecast Configuration\n",
|
||||||
|
"# @markdown - **`forecast_init_time`**: The starting time for the forecast in ISO 8601 format (e.g., `2025-09-21T00:00:00Z`). Models are available for dates from 2024 onwards.\n",
|
||||||
|
"# @markdown - **`horizon_hrs`**: The desired length of the forecast in hours (e.g., 240 for a 10-day forecast).\n",
|
||||||
|
"# @markdown - **`model_seed`**: Choose a specific model seed (1-4) or select \"all\" to run inference with all four seeds in parallel for improved accuracy.\n",
|
||||||
|
"# @markdown - **`enable_hourly_prediction`**: If checked, the model will generate 1-hour predictions.\n",
|
||||||
|
"forecast_init_time = \"2025-11-20T00:00:00Z\" # @param {type:\"string\"}\n",
|
||||||
|
"horizon_hrs = 72 # @param {type:\"integer\"}\n",
|
||||||
|
"model_seed = \"all\" # @param [\"1\", \"2\", \"3\", \"4\", \"all\"]\n",
|
||||||
|
"enable_hourly_prediction = True # @param {type:\"boolean\"}\n",
|
||||||
|
"\n",
|
||||||
|
"# --- Parameter Validation and Configuration ---\n",
|
||||||
|
"\n",
|
||||||
|
"# Derive accelerator type and count from the chosen machine type\n",
|
||||||
|
"if machine_type.startswith(\"a3-highgpu\"):\n",
|
||||||
|
" accelerator_type = \"NVIDIA_H100_80GB\"\n",
|
||||||
|
"elif machine_type.startswith(\"a2-ultragpu\"):\n",
|
||||||
|
" accelerator_type = \"NVIDIA_A100_80GB\"\n",
|
||||||
|
"else:\n",
|
||||||
|
" raise ValueError(f\"Invalid machine type selected: {machine_type}.\")\n",
|
||||||
|
"\n",
|
||||||
|
"try:\n",
|
||||||
|
" # Extract the number of GPUs from the machine type string, e.g., 'a3-highgpu-4g' -> 4\n",
|
||||||
|
" accelerators_per_machine = int(re.search(r\"-(\\d+)g$\", machine_type).group(1))\n",
|
||||||
|
"except (AttributeError, ValueError):\n",
|
||||||
|
" raise ValueError(\n",
|
||||||
|
" f\"Could not determine accelerator count from machine type: {machine_type}\"\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
"seeds_to_run = [1, 2, 3, 4] if model_seed == \"all\" else [int(model_seed)]\n",
|
||||||
|
"num_seeds_to_run = len(seeds_to_run)\n",
|
||||||
|
"\n",
|
||||||
|
"num_samples_per_seed = num_samples\n",
|
||||||
|
"if len(seeds_to_run) > 1:\n",
|
||||||
|
" if num_samples % num_seeds_to_run != 0:\n",
|
||||||
|
" raise ValueError(\n",
|
||||||
|
" f\"`num_samples` ({num_samples}) is not divisible by the number of seeds to run ({num_seeds_to_run}.\"\n",
|
||||||
|
" )\n",
|
||||||
|
" num_samples_per_seed = num_samples // num_seeds_to_run\n",
|
||||||
|
"\n",
|
||||||
|
"# Validate that num_samples is a multiple of accelerators_per_machine\n",
|
||||||
|
"if num_samples_per_seed % accelerators_per_machine != 0:\n",
|
||||||
|
" raise ValueError(\n",
|
||||||
|
" f\"`num_samples_per_seed` ({num_samples_per_seed}) must be a multiple of the GPUs per machine ({accelerators_per_machine} for {machine_type}).\"\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
"# Calculate the number of replicas per seed\n",
|
||||||
|
"replica_count_per_seed = num_samples_per_seed // accelerators_per_machine\n",
|
||||||
|
"\n",
|
||||||
|
"# Ensure that there are enough samples\n",
|
||||||
|
"if num_samples_per_seed % accelerators_per_machine != 0:\n",
|
||||||
|
" raise ValueError(\n",
|
||||||
|
" f\"`num_samples_per_seed` ({num_samples_per_seed}) must be a multiple of the GPUs per machine ({accelerators_per_machine} for {machine_type}).\"\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
"# Calculate total GPUs needed for all jobs\n",
|
||||||
|
"total_gpus_needed = num_samples * (4 if model_seed == \"all\" else 1)\n",
|
||||||
|
"\n",
|
||||||
|
"# Set Docker URI\n",
|
||||||
|
"WEATHERNEXT2_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/weather-next-2-inference.gpu.0-1:latest\"\n",
|
||||||
|
"\n",
|
||||||
|
"print(\"--- Job Configuration Summary ---\")\n",
|
||||||
|
"print(f\"Total Samples: {num_samples}\")\n",
|
||||||
|
"print(f\"Machine Type: {machine_type}\")\n",
|
||||||
|
"print(f\"Accelerator Type: {accelerator_type}\")\n",
|
||||||
|
"print(f\"GPUs per Machine: {accelerators_per_machine}\")\n",
|
||||||
|
"print(f\"Total number seeds to run: {num_seeds_to_run}\")\n",
|
||||||
|
"print(f\"Total number samples per seed: {num_samples_per_seed}\")\n",
|
||||||
|
"print(f\"Calculated Machine Replicas Per Seed: {replica_count_per_seed}\")\n",
|
||||||
|
"print(f\"Total GPUs per Job: {num_samples}\")\n",
|
||||||
|
"print(f\"Total GPUs across all Jobs (ensure sufficient quota): {total_gpus_needed}\")\n",
|
||||||
|
"print(f\"Docker Image: {WEATHERNEXT2_DOCKER_URI}\")\n",
|
||||||
|
"print(\"---------------------------------\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {
|
||||||
|
"id": "VSuGx2jmpwP0"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# @title Run Forecasts\n",
|
||||||
|
"# @markdown This section creates and runs one or more Vertex AI Custom Training Jobs to generate the forecasts.\n",
|
||||||
|
"# @markdown **This operation is asynchronous.** The jobs will be submitted and this cell will complete quickly.\n",
|
||||||
|
"# @markdown You must monitor the job progress in the Google Cloud Console (https://console.cloud.google.com/vertex-ai/training/custom-jobs).\n",
|
||||||
|
"\n",
|
||||||
|
"from google.cloud.aiplatform.compat.types import \\\n",
|
||||||
|
" custom_job as gca_custom_job_compat\n",
|
||||||
|
"\n",
|
||||||
|
"print(f\"Submitting {len(seeds_to_run)} job(s) to run in parallel.\")\n",
|
||||||
|
"\n",
|
||||||
|
"launched_jobs = []\n",
|
||||||
|
"output_dirs = {}\n",
|
||||||
|
"\n",
|
||||||
|
"if scheduling_strategy == \"FLEX_START\":\n",
|
||||||
|
" SCHEDULLING_STRATEGY = gca_custom_job_compat.Scheduling.Strategy.FLEX_START\n",
|
||||||
|
"elif scheduling_strategy == \"SPOT\":\n",
|
||||||
|
" SCHEDULLING_STRATEGY = gca_custom_job_compat.Scheduling.Strategy.SPOT\n",
|
||||||
|
"else:\n",
|
||||||
|
" SCHEDULLING_STRATEGY = gca_custom_job_compat.Scheduling.Strategy.STANDARD\n",
|
||||||
|
"\n",
|
||||||
|
"for seed in seeds_to_run:\n",
|
||||||
|
" output_dir = os.path.join(BUCKET_URI, \"weathernext2_outputs\")\n",
|
||||||
|
" output_dirs[seed] = output_dir\n",
|
||||||
|
"\n",
|
||||||
|
" docker_args_list = [\n",
|
||||||
|
" f\"--pred_root_dir={output_dir}\",\n",
|
||||||
|
" f\"--num_samples={num_samples_per_seed}\",\n",
|
||||||
|
" f\"--horizon_hrs={horizon_hrs}\",\n",
|
||||||
|
" f\"--forecast_init_time={forecast_init_time}\",\n",
|
||||||
|
" f\"--model_seed={seed}\",\n",
|
||||||
|
" f\"--enable_hourly_prediction={enable_hourly_prediction}\",\n",
|
||||||
|
" ]\n",
|
||||||
|
"\n",
|
||||||
|
" JOB_NAME = get_job_name_with_datetime(\n",
|
||||||
|
" prefix=f\"wn2-forecast-s{seed}-n{num_samples_per_seed}\"\n",
|
||||||
|
" )\n",
|
||||||
|
" print(f\"\\n--- Submitting Job for Seed {seed} ---\")\n",
|
||||||
|
" print(f\"JOB_NAME: {JOB_NAME}\")\n",
|
||||||
|
"\n",
|
||||||
|
" job = aiplatform.CustomContainerTrainingJob(\n",
|
||||||
|
" display_name=JOB_NAME,\n",
|
||||||
|
" container_uri=WEATHERNEXT2_DOCKER_URI,\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" job.run(\n",
|
||||||
|
" args=docker_args_list,\n",
|
||||||
|
" replica_count=replica_count_per_seed,\n",
|
||||||
|
" machine_type=machine_type,\n",
|
||||||
|
" accelerator_type=accelerator_type,\n",
|
||||||
|
" accelerator_count=accelerators_per_machine,\n",
|
||||||
|
" scheduling_strategy=SCHEDULLING_STRATEGY,\n",
|
||||||
|
" # Change this to True if you need to debug why the job hasn't started\n",
|
||||||
|
" sync=False,\n",
|
||||||
|
" )\n",
|
||||||
|
" launched_jobs.append(job)\n",
|
||||||
|
" print(\n",
|
||||||
|
" \"--> Job submitted successfully. Monitor it in the Google Cloud Console at https://console.cloud.google.com/vertex-ai/training/custom-jobs\"\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
"print(\"\\nAll forecast jobs have been submitted.\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {
|
||||||
|
"cellView": "form",
|
||||||
|
"id": "x6ZVqWRopWcI"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# @title Visualize Forecasts (Unified)\n",
|
||||||
|
"# @markdown Select which forecast output you want to visualize. This single component\n",
|
||||||
|
"# @markdown can handle both the standard 6-hourly predictions and the datasets\n",
|
||||||
|
"# @markdown with 1-hour model (which have a 'subtime' dimension).\n",
|
||||||
|
"# @markdown If you run into `Error loading Zarr store: unrecognized engine 'zarr'...` try restarting the runtime session and reruning this cell.\n",
|
||||||
|
"\n",
|
||||||
|
"# @markdown ---\n",
|
||||||
|
"# @markdown ### Visualization Settings\n",
|
||||||
|
"# @markdown - **`model_seed_to_visualize`**: Choose a specific model seed (1-4) to visualize. This should be one of the model seeds selected in the **Forecast Configuration** above.\n",
|
||||||
|
"# @markdown - **`time_steps_to_visualize`**: Choose to visualize 1-hourly or 6-hourly forecasts. If 1-hourly is selected, ensure `enable_hourly_prediction` was selected in the **Forecast Configuration** above.\n",
|
||||||
|
"# @markdown - **`variable_to_visualize`**: Choose the weather variable to visualize. See the [WeatherNext documentation](https://developers.google.com/weathernext/guides/model-specs-vmg) for variable names and descriptions.\n",
|
||||||
|
"# @markdown - **`sample_to_visualize`**: Choose the sample (ensemble member) to visualize.\n",
|
||||||
|
"# @markdown - **`plot_size`**: Choose the size of the plot to generate.\n",
|
||||||
|
"model_seed_to_visualize = \"4\" # @param [\"1\", \"2\", \"3\", \"4\"]\n",
|
||||||
|
"time_steps_to_visualize = \"6-Hourly\" # @param [\"6-Hourly\", \"1-Hourly\"]\n",
|
||||||
|
"variable_to_visualize = \"2m_temperature\" # @param {type:\"string\"}\n",
|
||||||
|
"sample_to_visualize = 0 # @param {type:\"integer\"}\n",
|
||||||
|
"plot_size = 8 # @param {type:\"number\"}\n",
|
||||||
|
"level_to_visualize = None\n",
|
||||||
|
"# @markdown ---\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# --- Helper Functions ---\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def init_time_to_folder_path(init_time: str) -> str:\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" Convert init time to expected GCS folder path.\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" init_date, init_time = init_time.split(\"T\")\n",
|
||||||
|
" return f\"{init_date.replace('-', '')}_{init_time[0:2]}hr\"\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# override these if you'd like to visualize a different set of forecasts\n",
|
||||||
|
"visualize_bucket = BUCKET_URI\n",
|
||||||
|
"visualize_init_date = forecast_init_time\n",
|
||||||
|
"\n",
|
||||||
|
"# set paths based on chosen model seed, bucket, and init date\n",
|
||||||
|
"path_to_6hr_zarr = f\"{BUCKET_URI}/weathernext2_outputs/weathernext_2_seed_{model_seed_to_visualize}/{init_time_to_folder_path(visualize_init_date)}_01_preds/predictions.zarr/\"\n",
|
||||||
|
"path_to_1hr_zarr = f\"{BUCKET_URI}/weathernext2_outputs/weathernext_2_seed_{model_seed_to_visualize}_hourly/{init_time_to_folder_path(visualize_init_date)}_01_preds/predictions.zarr/\"\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"import datetime\n",
|
||||||
|
"from typing import Optional\n",
|
||||||
|
"\n",
|
||||||
|
"import matplotlib\n",
|
||||||
|
"import matplotlib.animation as animation\n",
|
||||||
|
"import matplotlib.pyplot as plt\n",
|
||||||
|
"import numpy as np\n",
|
||||||
|
"import xarray\n",
|
||||||
|
"from IPython.display import HTML\n",
|
||||||
|
"\n",
|
||||||
|
"matplotlib.rcParams[\"animation.embed_limit\"] = 500\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# --- Helper Functions ---\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def select_data(\n",
|
||||||
|
" data: xarray.Dataset,\n",
|
||||||
|
" variable: str,\n",
|
||||||
|
" level: Optional[int] = None,\n",
|
||||||
|
") -> xarray.Dataset:\n",
|
||||||
|
" \"\"\"Selects a variable from the dataset and optionally a level.\"\"\"\n",
|
||||||
|
" data = data[variable]\n",
|
||||||
|
" if \"batch\" in data.dims:\n",
|
||||||
|
" data = data.isel(batch=0)\n",
|
||||||
|
" if level is not None and \"level\" in data.coords:\n",
|
||||||
|
" data = data.sel(level=level)\n",
|
||||||
|
" return data\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def scale_data(\n",
|
||||||
|
" data: xarray.Dataset,\n",
|
||||||
|
" center: Optional[float] = None,\n",
|
||||||
|
" robust: bool = False,\n",
|
||||||
|
") -> tuple[xarray.Dataset, matplotlib.colors.Normalize, str]:\n",
|
||||||
|
" \"\"\"Scales the data for visualization.\"\"\"\n",
|
||||||
|
" vmin = np.nanpercentile(data.values, (2 if robust else 0))\n",
|
||||||
|
" vmax = np.nanpercentile(data.values, (98 if robust else 100))\n",
|
||||||
|
" if center is not None:\n",
|
||||||
|
" diff = max(vmax - center, center - vmin)\n",
|
||||||
|
" vmin = center - diff\n",
|
||||||
|
" vmax = center + diff\n",
|
||||||
|
" return (\n",
|
||||||
|
" data,\n",
|
||||||
|
" matplotlib.colors.Normalize(vmin, vmax),\n",
|
||||||
|
" (\"RdBu_r\" if center is not None else \"viridis\"),\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def create_forecast_animation(\n",
|
||||||
|
" dataset: xarray.Dataset,\n",
|
||||||
|
" fig_title: str,\n",
|
||||||
|
" plot_size: float = 5,\n",
|
||||||
|
" robust: bool = False,\n",
|
||||||
|
") -> HTML:\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" Creates a forecast animation from an xarray Dataset.\n",
|
||||||
|
" It intelligently handles datasets with or without a 'subtime' dimension.\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" # --- Data Preparation ---\n",
|
||||||
|
" # Check if the data still has 'subtime'). If so, stack dimensions.\n",
|
||||||
|
" # Otherwise, just rename the 'time' dimension for consistency.\n",
|
||||||
|
" if \"subtime\" in dataset.dims:\n",
|
||||||
|
" print(\"Detected 'subtime' dimension. Stacking for hourly animation.\")\n",
|
||||||
|
" # Stack 'time' and 'subtime' into a single animation dimension\n",
|
||||||
|
" plot_data = dataset.stack(animation_step=(\"time\", \"subtime\")).transpose(\n",
|
||||||
|
" \"animation_step\", \"lat\", \"lon\"\n",
|
||||||
|
" )\n",
|
||||||
|
" else:\n",
|
||||||
|
" print(\"No 'subtime' dimension found. Using 'time' for 6-hourly animation.\")\n",
|
||||||
|
" # Use 'time' as the animation dimension\n",
|
||||||
|
" plot_data = dataset.rename({\"time\": \"animation_step\"})\n",
|
||||||
|
"\n",
|
||||||
|
" # Now, the animation dimension is always called 'animation_step'\n",
|
||||||
|
" max_steps = plot_data.sizes[\"animation_step\"]\n",
|
||||||
|
" init_time = plot_data.coords[\"init_time\"].values\n",
|
||||||
|
"\n",
|
||||||
|
" # Scale the data for color mapping\n",
|
||||||
|
" scaled_data, norm, cmap = scale_data(plot_data, robust=robust)\n",
|
||||||
|
"\n",
|
||||||
|
" # --- Plotting Setup ---\n",
|
||||||
|
" figure = plt.figure(figsize=(plot_size * 2, plot_size))\n",
|
||||||
|
" ax = figure.add_subplot(1, 1, 1)\n",
|
||||||
|
" ax.set_xticks([])\n",
|
||||||
|
" ax.set_yticks([])\n",
|
||||||
|
" figure.suptitle(fig_title, fontsize=16)\n",
|
||||||
|
" figure.tight_layout(rect=[0, 0.03, 1, 0.95]) # Adjust for title\n",
|
||||||
|
"\n",
|
||||||
|
" im = ax.imshow(\n",
|
||||||
|
" scaled_data.isel(animation_step=0), norm=norm, origin=\"lower\", cmap=cmap\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" plt.colorbar(\n",
|
||||||
|
" mappable=im,\n",
|
||||||
|
" ax=ax,\n",
|
||||||
|
" orientation=\"vertical\",\n",
|
||||||
|
" pad=0.02,\n",
|
||||||
|
" aspect=16,\n",
|
||||||
|
" shrink=0.75,\n",
|
||||||
|
" cmap=cmap,\n",
|
||||||
|
" extend=(\"both\" if robust else \"neither\"),\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" # --- Animation Update Function ---\n",
|
||||||
|
" def update(frame):\n",
|
||||||
|
" # Get the coordinates for the current frame\n",
|
||||||
|
" step_coords = plot_data[\"animation_step\"][frame].coords\n",
|
||||||
|
"\n",
|
||||||
|
" # Calculate total offset and valid time based on available coordinates\n",
|
||||||
|
" if \"subtime\" in step_coords: # Hourly data\n",
|
||||||
|
" total_offset = step_coords[\"time\"].values + step_coords[\"subtime\"].values\n",
|
||||||
|
" else: # 6-hourly data\n",
|
||||||
|
" total_offset = step_coords[\"animation_step\"].values\n",
|
||||||
|
"\n",
|
||||||
|
" total_hours = total_offset / np.timedelta64(1, \"h\")\n",
|
||||||
|
" valid_time = init_time + total_offset\n",
|
||||||
|
" valid_time_str = np.datetime_as_string(valid_time, unit=\"m\").replace(\"T\", \" \")\n",
|
||||||
|
"\n",
|
||||||
|
" new_title = (\n",
|
||||||
|
" f\"{fig_title}\\n\"\n",
|
||||||
|
" f\"Valid: {valid_time_str} UTC (Forecast: +{total_hours:.1f}h)\"\n",
|
||||||
|
" )\n",
|
||||||
|
" figure.suptitle(new_title, fontsize=16)\n",
|
||||||
|
" im.set_data(scaled_data.isel(animation_step=frame))\n",
|
||||||
|
"\n",
|
||||||
|
" # --- Create and Display Animation ---\n",
|
||||||
|
" ani = animation.FuncAnimation(\n",
|
||||||
|
" fig=figure, func=update, frames=max_steps, interval=250\n",
|
||||||
|
" )\n",
|
||||||
|
" plt.close(figure.number)\n",
|
||||||
|
" return HTML(ani.to_html5_video())\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# --- Main Visualization Logic ---\n",
|
||||||
|
"\n",
|
||||||
|
"# 1. Select the correct path based on the user's dropdown choice\n",
|
||||||
|
"if time_steps_to_visualize == \"6-Hourly\":\n",
|
||||||
|
" path_to_zarr = path_to_6hr_zarr\n",
|
||||||
|
"elif time_steps_to_visualize == \"1-Hourly\":\n",
|
||||||
|
" path_to_zarr = path_to_1hr_zarr\n",
|
||||||
|
"else:\n",
|
||||||
|
" raise ValueError(\"Invalid visualization target selected.\")\n",
|
||||||
|
"\n",
|
||||||
|
"print(f\"Loading data from: {path_to_zarr}\")\n",
|
||||||
|
"\n",
|
||||||
|
"# 2. Load the dataset\n",
|
||||||
|
"try:\n",
|
||||||
|
" full_dataset = xarray.open_zarr(path_to_zarr)\n",
|
||||||
|
"except Exception as e:\n",
|
||||||
|
" print(f\"Error loading Zarr store: {e}\")\n",
|
||||||
|
" # This is a common point of failure, so we exit gracefully.\n",
|
||||||
|
"else:\n",
|
||||||
|
" # 3. Select the specific data slice for visualization\n",
|
||||||
|
" data_for_vis = full_dataset.isel(sample=sample_to_visualize)\n",
|
||||||
|
" variable_data = select_data(data_for_vis, variable_to_visualize, level_to_visualize)\n",
|
||||||
|
"\n",
|
||||||
|
" # 4. Generate the title\n",
|
||||||
|
" title = f\"{variable_to_visualize} (Sample {sample_to_visualize})\"\n",
|
||||||
|
" if level_to_visualize:\n",
|
||||||
|
" title += f\" at {level_to_visualize} hPa\"\n",
|
||||||
|
"\n",
|
||||||
|
" # 5. Create and display the animation\n",
|
||||||
|
" display(create_forecast_animation(variable_data, title, plot_size, robust=True))"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"metadata": {
|
||||||
|
"colab": {
|
||||||
|
"name": "weathernext_2_dws.ipynb",
|
||||||
|
"toc_visible": true
|
||||||
|
},
|
||||||
|
"kernelspec": {
|
||||||
|
"display_name": "Python 3",
|
||||||
|
"name": "python3"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"nbformat": 4,
|
||||||
|
"nbformat_minor": 0
|
||||||
|
}
|
||||||
@@ -0,0 +1,970 @@
|
|||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "1a737002",
|
||||||
|
"metadata": {
|
||||||
|
"id": "RirpY96_3zL0"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Copyright 2026 Google LLC\n",
|
||||||
|
"\n",
|
||||||
|
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
|
||||||
|
"# you may not use this file except in compliance with the License.\n",
|
||||||
|
"# You may obtain a copy of the License at\n",
|
||||||
|
"\n",
|
||||||
|
"# https://www.apache.org/licenses/LICENSE-2.0\n",
|
||||||
|
"\n",
|
||||||
|
"# Unless required by applicable law or agreed to in writing, software\n",
|
||||||
|
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
|
||||||
|
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
|
||||||
|
"# See the License for the specific language governing permissions and\n",
|
||||||
|
"# limitations under the License."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "bf59250d",
|
||||||
|
"metadata": {
|
||||||
|
"id": "52d0938f"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"# WeatherNext 2\n",
|
||||||
|
"<table align=\"left\">\n",
|
||||||
|
" <td style=\"text-align: center\">\n",
|
||||||
|
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/weathernext/weathernext_2_ic_pc.ipynb\">\n",
|
||||||
|
" <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
|
||||||
|
" </a>\n",
|
||||||
|
" </td>\n",
|
||||||
|
" <td style=\"text-align: center\">\n",
|
||||||
|
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fweathernext%2Fweathernext_2_ic_pc.ipynb\">\n",
|
||||||
|
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
|
||||||
|
" </a>\n",
|
||||||
|
" </td>\n",
|
||||||
|
" <td style=\"text-align: center\">\n",
|
||||||
|
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/weathernext/weathernext_2_ic_pc.ipynb\">\n",
|
||||||
|
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Enterprise Gemini Agent Platform logo\"><br> Open in Enterprise Gemini Agent Platform Workbench\n",
|
||||||
|
" </a>\n",
|
||||||
|
" </td>\n",
|
||||||
|
" <td style=\"text-align: center\">\n",
|
||||||
|
" <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/weathernext/weathernext_2_ic_pc.ipynb\">\n",
|
||||||
|
" <img width=\"32px\"src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
|
||||||
|
" </a>\n",
|
||||||
|
" </td>\n",
|
||||||
|
"</table>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "3c096b65",
|
||||||
|
"metadata": {
|
||||||
|
"id": "DGFzy9MTuYL_"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"## Overview\n",
|
||||||
|
"\n",
|
||||||
|
"This notebook demonstrates running [WeatherNext 2 inference on Google Cloud Enterprise Gemini Agent Platform](https://developers.google.com/weathernext/guides/access-vmg) using **customer-provided initial conditions** (custom inputs). WeatherNext 2 is Google's latest medium-range probabilistic forecasting model, principally an operational version the FGN model ([published June 2025](https://arxiv.org/abs/2506.10772)). More information is available in the [WeatherNext documentation](https://developers.google.com/weathernext).\n",
|
||||||
|
"\n",
|
||||||
|
"WeatherNext 2 supports customer-provided initial conditions for inference. Instead of using the default ECMWF HRES real-time data, customers can supply their own input Zarr files (e.g., from GFS or their own analysis systems) to generate forecasts with WeatherNext models.\n",
|
||||||
|
"\n",
|
||||||
|
"> **Important**: The model is **not fine-tuned** on custom input data. Forecast performance when using custom inputs is **not guaranteed** to match the quality achieved with the default ECMWF HRES inputs. Customers should perform their own evaluation of output quality.\n",
|
||||||
|
"\n",
|
||||||
|
"### Objective\n",
|
||||||
|
"\n",
|
||||||
|
"- Configure the model inputs for distributed, multi-host inference on H100 or A100 GPUs.\n",
|
||||||
|
"- Provide custom initial conditions (Zarr files) for model inference.\n",
|
||||||
|
"- Run WeatherNext 2 model forecasts in parallel.\n",
|
||||||
|
"- Visualize forecast results.\n",
|
||||||
|
"\n",
|
||||||
|
"### Costs\n",
|
||||||
|
"\n",
|
||||||
|
"This uses billable components of Google Cloud:\n",
|
||||||
|
"\n",
|
||||||
|
"* [Gemini Enterprise Agent Platform]( https://docs.cloud.google.com/gemini-enterprise-agent-platform)\n",
|
||||||
|
"* [Cloud Storage](https://cloud.google.com/storage/docs)\n",
|
||||||
|
"* [Gemini Enterprise Agent Platform Persistent Resource](https://docs.cloud.google.com/gemini-enterprise-agent-platform/machine-learning/training/persistent-resource-create)\n",
|
||||||
|
"\n",
|
||||||
|
"Learn about [Gemini Enterprise Agent Platform pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), [Gemini Enterprise Platform Persistent Resource](https://cloud.google.com/products/gemini-enterprise-agent-platform/pricing#custom-trained-models) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage.\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"## Before you begin\n",
|
||||||
|
"\n",
|
||||||
|
"### Request For GPU Quota\n",
|
||||||
|
"\n",
|
||||||
|
"**WARNING:** Make sure you have sufficient GPU quota allocated for the inference configuration (i.e. `num_samples`) before running Gemini Enterprise Agent Platform Jobs or provisioning a Persistent Resource. Otherwise, some Gemini Enterprise Agent Platform jobs may run while others will fail which would produce\n",
|
||||||
|
"incomplete results.\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"By default, the quota for GPUs is 0. You can request a higher quota by following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
|
||||||
|
"\n",
|
||||||
|
"You will need to request quota for either **NVIDIA H100 80GB GPUs** or **NVIDIA A100 80GB GPUs** in your selected region. The total number of GPUs you request must be sufficient for your largest planned forecast (i.e., `num_samples`) or your provisioned Persistent Resource capacity.\n",
|
||||||
|
"\n",
|
||||||
|
"Depending on whether you use a **Persistent Resource** or run standard custom jobs, you should request the following quota under Service `Enterprise Gemini Agent Platform API`:\n",
|
||||||
|
"\n",
|
||||||
|
"#### Option 1: Persistent Resource\n",
|
||||||
|
"- Name: `Persistent resource Nvidia A100 80GB GPUs per region` OR `Persistent resource Nvidia H100 GPUs per region`\n",
|
||||||
|
"\n",
|
||||||
|
"#### Option 2: Standard Custom Model Training (Preemptible)\n",
|
||||||
|
"- Name: `Custom model training preemptible Nvidia A100 80GB GPUs per region` OR `Custom model training preemptible Nvidia H100 GPUs per region`\n",
|
||||||
|
"\n",
|
||||||
|
"### Custom Inputs Guide\n",
|
||||||
|
"\n",
|
||||||
|
"```\n",
|
||||||
|
"The customer provides input Zarr files in a GCS bucket path via the\n",
|
||||||
|
"`--cns_data_dir` flag. \n",
|
||||||
|
"\n",
|
||||||
|
"#### How to Use\n",
|
||||||
|
"\n",
|
||||||
|
"##### Flag: `--input_data_gcs_dir`\n",
|
||||||
|
"\n",
|
||||||
|
"`--input_data_gcs_dir` flag for specifying custom input data:\n",
|
||||||
|
"\n",
|
||||||
|
"```\n",
|
||||||
|
"--input_data_gcs_dir=gs://customer-bucket/path-to-input-data\n",
|
||||||
|
"```\n",
|
||||||
|
"\n",
|
||||||
|
"##### GCS Bucket Setup\n",
|
||||||
|
"\n",
|
||||||
|
"We recommend using the **same GCS bucket** for both input data and output\n",
|
||||||
|
"predictions.\n",
|
||||||
|
"\n",
|
||||||
|
"The customer should place their input Zarr files under a path within their\n",
|
||||||
|
"existing output bucket, e.g.:\n",
|
||||||
|
"\n",
|
||||||
|
"```\n",
|
||||||
|
"gs://customer-bucket/custom-inputs/ ← input Zarr files go here\n",
|
||||||
|
"gs://customer-bucket/outputs/ ← model predictions are written here\n",
|
||||||
|
"```\n",
|
||||||
|
"\n",
|
||||||
|
"#### Input File Format\n",
|
||||||
|
"\n",
|
||||||
|
"##### Zarr V3 Format Requirement\n",
|
||||||
|
"\n",
|
||||||
|
"**All custom input Zarr files MUST be in Zarr V3 format.**\n",
|
||||||
|
"\n",
|
||||||
|
"When creating custom input files, ensure they are saved in Zarr V3 format:\n",
|
||||||
|
"\n",
|
||||||
|
"```python\n",
|
||||||
|
"import xarray as xa\n",
|
||||||
|
"\n",
|
||||||
|
"# When creating input files, ensure they are saved in Zarr V3 format\n",
|
||||||
|
"dataset.to_zarr(\"path/to/output.zarr\", zarr_format=3)\n",
|
||||||
|
"```\n",
|
||||||
|
"\n",
|
||||||
|
"If V2 format files are provided as custom input, the job raises\n",
|
||||||
|
"an error when attempting to read them.\n",
|
||||||
|
"\n",
|
||||||
|
"##### Example Dataset\n",
|
||||||
|
"\n",
|
||||||
|
"Opening an example input Zarr file with xarray (init time 2026-03-18 00:00 UTC):\n",
|
||||||
|
"\n",
|
||||||
|
"```python\n",
|
||||||
|
">>> import xarray as xa\n",
|
||||||
|
">>> ds = xa.open_zarr(\"2026_A1D03180000031800011.zarr\")\n",
|
||||||
|
">>> ds\n",
|
||||||
|
"<xarray.Dataset> Size: ...\n",
|
||||||
|
"Dimensions: (isobaricInhPa: 13, latitude: 721, longitude: 1440)\n",
|
||||||
|
"Coordinates:\n",
|
||||||
|
" * isobaricInhPa (isobaricInhPa) float64 13 ...\n",
|
||||||
|
" * latitude (latitude) float64 721 ...\n",
|
||||||
|
" * longitude (longitude) float64 1440 ...\n",
|
||||||
|
" number int64 ...\n",
|
||||||
|
" step int64 ...\n",
|
||||||
|
" surface float64 ...\n",
|
||||||
|
" time int64 ...\n",
|
||||||
|
" valid_time int64 ...\n",
|
||||||
|
"Data variables: (13 total)\n",
|
||||||
|
" msl (latitude, longitude) float32 ...\n",
|
||||||
|
" q (isobaricInhPa, latitude, longitude) float32 ...\n",
|
||||||
|
" sst (latitude, longitude) float32 ...\n",
|
||||||
|
" t (isobaricInhPa, latitude, longitude) float32 ...\n",
|
||||||
|
" t2m (latitude, longitude) float32 ...\n",
|
||||||
|
" u (isobaricInhPa, latitude, longitude) float32 ...\n",
|
||||||
|
" u10 (latitude, longitude) float32 ...\n",
|
||||||
|
" u100 (latitude, longitude) float32 ...\n",
|
||||||
|
" v (isobaricInhPa, latitude, longitude) float32 ...\n",
|
||||||
|
" v10 (latitude, longitude) float32 ...\n",
|
||||||
|
" v100 (latitude, longitude) float32 ...\n",
|
||||||
|
" w (isobaricInhPa, latitude, longitude) float32 ...\n",
|
||||||
|
" z (isobaricInhPa, latitude, longitude) float32 ...\n",
|
||||||
|
"```\n",
|
||||||
|
"\n",
|
||||||
|
"##### Variables\n",
|
||||||
|
"\n",
|
||||||
|
"All data variables are `float32`.\n",
|
||||||
|
"\n",
|
||||||
|
"Only the variables and levels needed for inference are listed here. Your zarr\n",
|
||||||
|
"may contain additional variables and levels.\n",
|
||||||
|
"\n",
|
||||||
|
"###### Pressure Level Variables (3D: `isobaricInhPa` × `latitude` × `longitude`)\n",
|
||||||
|
"\n",
|
||||||
|
"Shape: `[13, 721, 1440]`\n",
|
||||||
|
"\n",
|
||||||
|
"These variables are used for input at the following pressure levels (hPa):\n",
|
||||||
|
"`50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000`\n",
|
||||||
|
"\n",
|
||||||
|
"| Short Name | Description |\n",
|
||||||
|
"| :--- | :--- |\n",
|
||||||
|
"| `q` | specific humidity |\n",
|
||||||
|
"| `t` | temperature |\n",
|
||||||
|
"| `u` | u component of wind |\n",
|
||||||
|
"| `v` | v component of wind |\n",
|
||||||
|
"| `w` | vertical velocity |\n",
|
||||||
|
"| `z` | geopotential |\n",
|
||||||
|
"\n",
|
||||||
|
"###### Surface Variables (2D: `latitude` × `longitude`)\n",
|
||||||
|
"\n",
|
||||||
|
"Shape: `[721, 1440]`\n",
|
||||||
|
"\n",
|
||||||
|
"| Short Name | Description |\n",
|
||||||
|
"| :--- | :--- |\n",
|
||||||
|
"| `msl` | mean sea level pressure |\n",
|
||||||
|
"| `sst` | sea surface temperature |\n",
|
||||||
|
"| `t2m` | 2m temperature |\n",
|
||||||
|
"| `u10` | 10m u component of wind |\n",
|
||||||
|
"| `u100` | 100m u component of wind |\n",
|
||||||
|
"| `v10` | 10m v component of wind |\n",
|
||||||
|
"| `v100` | 100m v component of wind |\n",
|
||||||
|
"\n",
|
||||||
|
"###### Scalar Coordinates\n",
|
||||||
|
"\n",
|
||||||
|
"| Name | dtype | Description |\n",
|
||||||
|
"|--------------|-----------|----------------------------------------------|\n",
|
||||||
|
"| `number` | `int64` | Ensemble member number |\n",
|
||||||
|
"| `step` | `int64` | Forecast step |\n",
|
||||||
|
"| `surface` | `float64` | Surface level indicator |\n",
|
||||||
|
"| `time` | `int64` | Timestamp (units: days since init time, calendar: proleptic_gregorian) |\n",
|
||||||
|
"| `valid_time` | `int64` | Validity time |\n",
|
||||||
|
"\n",
|
||||||
|
"##### Dimension Coordinates\n",
|
||||||
|
"\n",
|
||||||
|
"| Coordinate | dtype | Shape | Description |\n",
|
||||||
|
"|-----------------|-----------|----------|--------------------------------|\n",
|
||||||
|
"| `isobaricInhPa` | `float64` | `[13]` | 13 pressure levels in hPa |\n",
|
||||||
|
"| `latitude` | `float64` | `[721]` | 0.25° resolution, 721 points |\n",
|
||||||
|
"| `longitude` | `float64` | `[1440]` | 0.25° resolution, 1440 points |\n",
|
||||||
|
"\n",
|
||||||
|
"##### Spatial Resolution\n",
|
||||||
|
"\n",
|
||||||
|
"The data is at **0.25° resolution** globally:\n",
|
||||||
|
"\n",
|
||||||
|
"- Latitude: 721 points (90°N to 90°S)\n",
|
||||||
|
"- Longitude: 1440 points (0° to 359.75°E)\n",
|
||||||
|
"\n",
|
||||||
|
"#### File Naming and Structure\n",
|
||||||
|
"\n",
|
||||||
|
"##### File Naming Convention\n",
|
||||||
|
"\n",
|
||||||
|
"The inference binary expects input Zarr files to follow a specific file naming\n",
|
||||||
|
"convention. Each file corresponds to a specific forecast initialization time and\n",
|
||||||
|
"uses the following format:\n",
|
||||||
|
"\n",
|
||||||
|
"```\n",
|
||||||
|
"<year>_<config><stream><MMDDHHMMMMDDHHMMEE>.zarr\n",
|
||||||
|
"```\n",
|
||||||
|
"\n",
|
||||||
|
"Where:\n",
|
||||||
|
"\n",
|
||||||
|
"- `<year>`: 4-digit year (e.g., `2025`)\n",
|
||||||
|
"- `<config>`: Data config name (default: `A1`)\n",
|
||||||
|
"- `<stream>`: `D` for 00/12 UTC init times, `S` for 06/18 UTC\n",
|
||||||
|
"- First `MMDDHHMM`: Month, day, hour, minute of the forecast init time\n",
|
||||||
|
"- Second `MMDDHHMM`: Month, day, hour, minute of the validity time\n",
|
||||||
|
"- `EE`: Experiment version (default: `1`)\n",
|
||||||
|
"\n",
|
||||||
|
"The validity minute is hardcoded to `01`.\n",
|
||||||
|
"\n",
|
||||||
|
"###### Examples\n",
|
||||||
|
"\n",
|
||||||
|
"For a forecast initialized at **2026-03-18 00:00 UTC** using fc0:\n",
|
||||||
|
"```\n",
|
||||||
|
"2026_A1D03180000031800011.zarr\n",
|
||||||
|
"```\n",
|
||||||
|
"\n",
|
||||||
|
"For a forecast initialized at **2026-03-17 06:00 UTC** using fc0:\n",
|
||||||
|
"```\n",
|
||||||
|
"2026_A1S03170600031706011.zarr\n",
|
||||||
|
"```\n",
|
||||||
|
"\n",
|
||||||
|
"##### Success Sentinels\n",
|
||||||
|
"\n",
|
||||||
|
"Each Zarr file directory must contain a `success` sentinel file to signal that\n",
|
||||||
|
"the data is complete and ready for reading:\n",
|
||||||
|
"\n",
|
||||||
|
"```\n",
|
||||||
|
"gs://<bucket>/custom-inputs/\n",
|
||||||
|
"├── 2026_A1D03180000031800011.zarr/\n",
|
||||||
|
"│ ├── .zmetadata\n",
|
||||||
|
"│ ├── <array data>\n",
|
||||||
|
"│ └── success ← required sentinel file\n",
|
||||||
|
"├── 2026_A1D03171200031712011.zarr/\n",
|
||||||
|
"│ ├── .zmetadata\n",
|
||||||
|
"│ ├── <array data>\n",
|
||||||
|
"│ └── success\n",
|
||||||
|
"```\n",
|
||||||
|
"\n",
|
||||||
|
"The sentinel is a zero-byte file named `success` placed inside each `.zarr`\n",
|
||||||
|
"directory. The job will not proceed with inference until all required input\n",
|
||||||
|
"sentinels exist.\n",
|
||||||
|
"\n",
|
||||||
|
"##### Number of Input Files\n",
|
||||||
|
"\n",
|
||||||
|
"The model typically requires **2 input timestamps** (the forecast init time and\n",
|
||||||
|
"6 hours prior). For example, for a forecast initialized at 2026-03-18 12:00 UTC,\n",
|
||||||
|
"the binary expects:\n",
|
||||||
|
"\n",
|
||||||
|
"1. `2026_A1D03181200031812011.zarr` (init time: 12:00 UTC)\n",
|
||||||
|
"2. `2026_A1D03180600031806011.zarr` (6 hours prior: 06:00 UTC)\n",
|
||||||
|
"\n",
|
||||||
|
"#### Caveats and Limitations\n",
|
||||||
|
"\n",
|
||||||
|
"1. **No fine-tuning guarantee**: The model is trained on ECMWF HRES\n",
|
||||||
|
" data. Using custom inputs from a different source may degrade\n",
|
||||||
|
" forecast quality, especially for features sensitive to the initial\n",
|
||||||
|
" condition source.\n",
|
||||||
|
"\n",
|
||||||
|
"2. **Variable completeness**: All variables listed above must be present in the\n",
|
||||||
|
" custom input files. Missing variables will cause the inference to fail.\n",
|
||||||
|
"\n",
|
||||||
|
"3. **Temporal alignment**: Custom input timestamps must align with valid HRES\n",
|
||||||
|
" forecast hours (00, 06, 12, or 18 UTC).\n",
|
||||||
|
"\n",
|
||||||
|
"4. **File format**: Only Zarr V3 format is accepted."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "d656ed2b",
|
||||||
|
"metadata": {
|
||||||
|
"id": "9349a5bece68"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"## Install packages - Restart the kernel after the installation"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "cbfab540",
|
||||||
|
"metadata": {
|
||||||
|
"id": "yt4GxkKcDD7Y"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# @title Install python packages\n",
|
||||||
|
"\n",
|
||||||
|
"# Note that you may need to restart the kernel after this step.\n",
|
||||||
|
"# If so, continue to the next cell after restarting.\n",
|
||||||
|
"\n",
|
||||||
|
"print(\"Installing python packages.\")\n",
|
||||||
|
"\n",
|
||||||
|
"! pip3 install \\\n",
|
||||||
|
" google-cloud-aiplatform==1.129.0 \\\n",
|
||||||
|
" xarray[complete]"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "c123b19f",
|
||||||
|
"metadata": {
|
||||||
|
"id": "48a20b368e8c"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"## Authenticate to Google Cloud Platform"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "3b7e5c45",
|
||||||
|
"metadata": {
|
||||||
|
"id": "pw5fTRvrDPQv"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"import sys\n",
|
||||||
|
"\n",
|
||||||
|
"if \"google.colab\" in sys.modules:\n",
|
||||||
|
" from google.colab import auth\n",
|
||||||
|
"\n",
|
||||||
|
" auth.authenticate_user()"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "7abff806",
|
||||||
|
"metadata": {
|
||||||
|
"id": "fd0ba250117b"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"## Set Google Cloud Project Pertinent Variables"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "b0805e58",
|
||||||
|
"metadata": {
|
||||||
|
"id": "iklx3Yb3Dmke"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Replace my_gcp_project with your gcp project\n",
|
||||||
|
"PROJECT_ID = \"<my_gcp_project>\" # @param {type:\"string\"}\n",
|
||||||
|
"\n",
|
||||||
|
"# Alternatively collect the default cloud project id from the OS env variable\n",
|
||||||
|
"# PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
|
||||||
|
"\n",
|
||||||
|
"# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
|
||||||
|
"# @markdown 2. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched.\n",
|
||||||
|
"\n",
|
||||||
|
"# Replace my_wn_bucket with your bucket\n",
|
||||||
|
"BUCKET_URI = \"gs://<my_wn_bucket>\" # @param {type:\"string\"}\n",
|
||||||
|
"\n",
|
||||||
|
"# @markdown 3. Select a region that has the required GPUs available.\n",
|
||||||
|
"\n",
|
||||||
|
"# Select a ***US*** region region\n",
|
||||||
|
"REGION = \"us-central1\" # @param {type:\"string\"}\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# @markdown 4. Provide the Persistence Resource ID\n",
|
||||||
|
"PERSISTENT_RESOURCE_ID = \"<my_persistent_resource_id>\" # @param {type:\"string\"}\n",
|
||||||
|
"\n",
|
||||||
|
"# Set Docker URI\n",
|
||||||
|
"# @markdown 5. Set the Image URI\n",
|
||||||
|
"WEATHERNEXT2_DOCKER_URI = \"us-central1-docker.pkg.dev/weathernext-1/wn25-private-preview-launch/weather-next-ic-2-inference.gpu.0-1:latest\""
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "9af9574a",
|
||||||
|
"metadata": {
|
||||||
|
"id": "KHbl4756DjMU"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"!gcloud config set project $PROJECT_ID"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "d914b484",
|
||||||
|
"metadata": {
|
||||||
|
"id": "dO2mF4CPfpHW1ZKbIsut8r69"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Import the necessary packages\n",
|
||||||
|
"import datetime\n",
|
||||||
|
"import os\n",
|
||||||
|
"import re\n",
|
||||||
|
"\n",
|
||||||
|
"from google.cloud import aiplatform\n",
|
||||||
|
"\n",
|
||||||
|
"# Enable the Enterprise Gemini Agent Platform API and Compute Engine API, if not already.\n",
|
||||||
|
"print(\"Enabling Enterprise Gemini Agent Platform API and Compute Engine API.\")\n",
|
||||||
|
"# ! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
|
||||||
|
"\n",
|
||||||
|
"# Cloud Storage bucket for storing the experiment artifacts.\n",
|
||||||
|
"now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
|
||||||
|
"BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
|
||||||
|
"\n",
|
||||||
|
"if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
|
||||||
|
" raise ValueError(\"GCS Bucket URI is invalid!\")\n",
|
||||||
|
"else:\n",
|
||||||
|
" assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
|
||||||
|
" shell_output = !gcloud storage buckets describe {BUCKET_NAME} --format=\"value(location)\" | tr '[:upper:]' '[:lower:]'\n",
|
||||||
|
" bucket_region = shell_output[0].strip().lower()\n",
|
||||||
|
" if bucket_region != REGION:\n",
|
||||||
|
" raise ValueError(\n",
|
||||||
|
" f\"Bucket region {bucket_region} is different from notebook region {REGION}\"\n",
|
||||||
|
" )\n",
|
||||||
|
"print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
|
||||||
|
"\n",
|
||||||
|
"STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
|
||||||
|
"# Initialize Enterprise Gemini Agent Platform API.\n",
|
||||||
|
"aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
|
||||||
|
"\n",
|
||||||
|
"# Utility functions\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def get_job_name_with_datetime(prefix: str) -> str:\n",
|
||||||
|
" return prefix + datetime.datetime.now().strftime(\"_%Y%m%d_%H%M%S\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "95c31932",
|
||||||
|
"metadata": {
|
||||||
|
"id": "cb667a3392f7"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"## Configure Model Parameters"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "f4854c78",
|
||||||
|
"metadata": {
|
||||||
|
"id": "22BW43yjps8D"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# @markdown ### Hardware Configuration for Distributed Inference\n",
|
||||||
|
"# @markdown - **`machine_type`**: Select a valid machine type. `a3-highgpu` series use NVIDIA H100 80GB GPUs. `a2-ultragpu` series use NVIDIA A100 80GB GPUs.\n",
|
||||||
|
"# @markdown - **`num_samples`**: The total number of ensemble members to generate.\n",
|
||||||
|
"# @markdown The number of machine replicas will be calculated automatically (`num_samples` / GPUs per machine). **Therefore, `num_samples` must be a multiple of the number of GPUs in your selected `machine_type`.**\n",
|
||||||
|
"machine_type = \"a2-ultragpu-1g\" # @param [\"a3-highgpu-1g\", \"a3-highgpu-2g\", \"a3-highgpu-4g\", \"a3-highgpu-8g\", \"a2-ultragpu-1g\", \"a2-ultragpu-2g\", \"a2-ultragpu-4g\", \"a2-ultragpu-8g\"]\n",
|
||||||
|
"num_samples = 4 # @param {type:\"integer\"}\n",
|
||||||
|
"\n",
|
||||||
|
"# @markdown ### Forecast Configuration\n",
|
||||||
|
"# @markdown - **`forecast_init_time`**: The starting time for the forecast in ISO 8601 format (e.g., `2025-09-21T00:00:00Z`). Models are available for dates from 2024 onwards.\n",
|
||||||
|
"# @markdown - **`horizon_hrs`**: The desired length of the forecast in hours (e.g., 240 for a 10-day forecast).\n",
|
||||||
|
"# @markdown - **`model_seed`**: Choose a specific model seed (1-4) or select \"all\" to run inference with all four seeds in parallel for improved accuracy.\n",
|
||||||
|
"# @markdown - **`enable_hourly_prediction`**: If checked, the model will generate 1-hour predictions.\n",
|
||||||
|
"forecast_init_time = \"2026-05-16T12:00:00Z\"\n",
|
||||||
|
"\n",
|
||||||
|
"horizon_hrs = 72 # @param {type:\"integer\"}\n",
|
||||||
|
"model_seed = \"all\" # @param [\"1\", \"2\", \"3\", \"4\", \"all\"]\n",
|
||||||
|
"enable_hourly_prediction = True # @param {type:\"boolean\"}\n",
|
||||||
|
"\n",
|
||||||
|
"# --- Parameter Validation and Configuration ---\n",
|
||||||
|
"\n",
|
||||||
|
"# Derive accelerator type and count from the chosen machine type\n",
|
||||||
|
"if machine_type.startswith(\"a3-highgpu\"):\n",
|
||||||
|
" accelerator_type = \"NVIDIA_H100_80GB\"\n",
|
||||||
|
"elif machine_type.startswith(\"a2-ultragpu\"):\n",
|
||||||
|
" accelerator_type = \"NVIDIA_A100_80GB\"\n",
|
||||||
|
"else:\n",
|
||||||
|
" raise ValueError(f\"Invalid machine type selected: {machine_type}.\")\n",
|
||||||
|
"\n",
|
||||||
|
"try:\n",
|
||||||
|
" # Extract the number of GPUs from the machine type string, e.g., 'a3-highgpu-4g' -> 4\n",
|
||||||
|
" accelerators_per_machine = int(re.search(r\"-(\\d+)g$\", machine_type).group(1))\n",
|
||||||
|
"except (AttributeError, ValueError):\n",
|
||||||
|
" raise ValueError(\n",
|
||||||
|
" f\"Could not determine accelerator count from machine type: {machine_type}\"\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
"seeds_to_run = [1, 2, 3, 4] if model_seed == \"all\" else [int(model_seed)]\n",
|
||||||
|
"num_seeds_to_run = len(seeds_to_run)\n",
|
||||||
|
"\n",
|
||||||
|
"num_samples_per_seed = num_samples\n",
|
||||||
|
"if len(seeds_to_run) > 1:\n",
|
||||||
|
" if num_samples % num_seeds_to_run != 0:\n",
|
||||||
|
" raise ValueError(\n",
|
||||||
|
" f\"`num_samples` ({num_samples}) is not divisible by the number of seeds to run ({num_seeds_to_run}.\"\n",
|
||||||
|
" )\n",
|
||||||
|
" num_samples_per_seed = num_samples // num_seeds_to_run\n",
|
||||||
|
"\n",
|
||||||
|
"# Validate that num_samples is a multiple of accelerators_per_machine\n",
|
||||||
|
"if num_samples_per_seed % accelerators_per_machine != 0:\n",
|
||||||
|
" raise ValueError(\n",
|
||||||
|
" f\"`num_samples_per_seed` ({num_samples_per_seed}) must be a multiple of the GPUs per machine ({accelerators_per_machine} for {machine_type}).\"\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
"# Calculate the number of replicas per seed\n",
|
||||||
|
"replica_count_per_seed = num_samples_per_seed // accelerators_per_machine\n",
|
||||||
|
"\n",
|
||||||
|
"# Ensure that there are enough samples\n",
|
||||||
|
"if num_samples_per_seed % accelerators_per_machine != 0:\n",
|
||||||
|
" raise ValueError(\n",
|
||||||
|
" f\"`num_samples_per_seed` ({num_samples_per_seed}) must be a multiple of the GPUs per machine ({accelerators_per_machine} for {machine_type}).\"\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
"# Calculate total GPUs needed for all jobs\n",
|
||||||
|
"total_gpus_needed = num_samples * (4 if model_seed == \"all\" else 1)\n",
|
||||||
|
"\n",
|
||||||
|
"print(\"--- Job Configuration Summary ---\")\n",
|
||||||
|
"print(f\"Total Samples: {num_samples}\")\n",
|
||||||
|
"print(f\"Machine Type: {machine_type}\")\n",
|
||||||
|
"print(f\"Accelerator Type: {accelerator_type}\")\n",
|
||||||
|
"print(f\"GPUs per Machine: {accelerators_per_machine}\")\n",
|
||||||
|
"print(f\"Total number seeds to run: {num_seeds_to_run}\")\n",
|
||||||
|
"print(f\"Total number samples per seed: {num_samples_per_seed}\")\n",
|
||||||
|
"print(f\"Calculated Machine Replicas Per Seed: {replica_count_per_seed}\")\n",
|
||||||
|
"print(f\"Total GPUs per Job: {num_samples}\")\n",
|
||||||
|
"print(f\"Total GPUs across all Jobs (ensure sufficient quota): {total_gpus_needed}\")\n",
|
||||||
|
"print(f\"Docker Image: {WEATHERNEXT2_DOCKER_URI}\")\n",
|
||||||
|
"print(\"---------------------------------\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "0ba7d73c",
|
||||||
|
"metadata": {
|
||||||
|
"id": "fe8b458ce060"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"## Run forecast jobs"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "42329e44",
|
||||||
|
"metadata": {
|
||||||
|
"id": "jVTDLIM87QSr"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# @markdown This section creates and runs one or more Enterprise Gemini Agent Platform Custom Training Jobs to generate the forecasts.\n",
|
||||||
|
"# @markdown **This operation is asynchronous.** The jobs will be submitted and this cell will complete quickly.\n",
|
||||||
|
"# @markdown You must monitor the job progress in the Google Cloud Console (https://console.cloud.google.com/vertex-ai/training/custom-jobs).\n",
|
||||||
|
"\n",
|
||||||
|
"import time\n",
|
||||||
|
"\n",
|
||||||
|
"print(\n",
|
||||||
|
" f\"Submitting {len(seeds_to_run)} job(s) to target persistent resource: {PERSISTENT_RESOURCE_ID}\"\n",
|
||||||
|
")\n",
|
||||||
|
"\n",
|
||||||
|
"launched_jobs = []\n",
|
||||||
|
"output_dirs = {}\n",
|
||||||
|
"\n",
|
||||||
|
"for seed in seeds_to_run:\n",
|
||||||
|
" output_dir = os.path.join(BUCKET_URI, \"weathernext2_outputs\")\n",
|
||||||
|
" output_dirs[seed] = output_dir\n",
|
||||||
|
"\n",
|
||||||
|
" docker_args_list = [\n",
|
||||||
|
" f\"--pred_root_dir={output_dir}\",\n",
|
||||||
|
" f\"--num_samples={num_samples_per_seed}\",\n",
|
||||||
|
" f\"--horizon_hrs={horizon_hrs}\",\n",
|
||||||
|
" f\"--forecast_init_time={forecast_init_time}\",\n",
|
||||||
|
" f\"--model_seed={seed}\",\n",
|
||||||
|
" f\"--enable_hourly_prediction={enable_hourly_prediction}\",\n",
|
||||||
|
" f\"--input_data_gcs_dir={BUCKET_URI}/custom_inputs/\",\n",
|
||||||
|
" ]\n",
|
||||||
|
"\n",
|
||||||
|
" # Define the worker pool spec for the custom job matching the warm pool specs\n",
|
||||||
|
" worker_pool_specs = [\n",
|
||||||
|
" {\n",
|
||||||
|
" \"machine_spec\": {\n",
|
||||||
|
" \"machine_type\": machine_type,\n",
|
||||||
|
" \"accelerator_type\": accelerator_type,\n",
|
||||||
|
" \"accelerator_count\": accelerators_per_machine,\n",
|
||||||
|
" },\n",
|
||||||
|
" \"disk_spec\": {\n",
|
||||||
|
" \"boot_disk_type\": \"pd-standard\",\n",
|
||||||
|
" \"boot_disk_size_gb\": 100,\n",
|
||||||
|
" },\n",
|
||||||
|
" \"replica_count\": 1,\n",
|
||||||
|
" \"container_spec\": {\n",
|
||||||
|
" \"image_uri\": WEATHERNEXT2_DOCKER_URI,\n",
|
||||||
|
" \"args\": docker_args_list,\n",
|
||||||
|
" },\n",
|
||||||
|
" }\n",
|
||||||
|
" ]\n",
|
||||||
|
"\n",
|
||||||
|
" JOB_NAME = get_job_name_with_datetime(prefix=f\"wn2-persistent-run-s{seed}\")\n",
|
||||||
|
" print(f\"\\n--- Submitting Job for Seed {seed} ---\")\n",
|
||||||
|
" print(f\"JOB_NAME: {JOB_NAME}\")\n",
|
||||||
|
"\n",
|
||||||
|
" custom_job = aiplatform.CustomJob(\n",
|
||||||
|
" display_name=JOB_NAME,\n",
|
||||||
|
" worker_pool_specs=worker_pool_specs,\n",
|
||||||
|
" project=PROJECT_ID,\n",
|
||||||
|
" location=REGION,\n",
|
||||||
|
" staging_bucket=f\"{BUCKET_URI}/custom_job_staging\",\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" # Run the job asynchronously targeting the warm persistent resource pool!\n",
|
||||||
|
" # Since the 4-GPU persistent resource was created without explicit custom SA runtime permissions,\n",
|
||||||
|
" # we run using the default Enterprise Gemini Agent Platform Service Agent (which already has GCS and registry reader permissions).\n",
|
||||||
|
" custom_job.run(\n",
|
||||||
|
" persistent_resource_id=PERSISTENT_RESOURCE_ID,\n",
|
||||||
|
" disable_retries=True,\n",
|
||||||
|
" sync=False,\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" # Wait for GCA resource generation details\n",
|
||||||
|
" print(\"Waiting for CustomJob details...\")\n",
|
||||||
|
" for _ in range(15):\n",
|
||||||
|
" try:\n",
|
||||||
|
" if custom_job.resource_name:\n",
|
||||||
|
" break\n",
|
||||||
|
" except Exception:\n",
|
||||||
|
" pass\n",
|
||||||
|
" time.sleep(1)\n",
|
||||||
|
"\n",
|
||||||
|
" job_id = custom_job.name.split(\"/\")[-1]\n",
|
||||||
|
" print(f\"--> Job submitted successfully. ID: {job_id}\")\n",
|
||||||
|
" print(\n",
|
||||||
|
" f\"Console Link: https://console.cloud.google.com/vertex-ai/locations/{REGION}/training/custom-jobs/{job_id}?project={PROJECT_ID}\"\n",
|
||||||
|
" )\n",
|
||||||
|
" launched_jobs.append(custom_job)\n",
|
||||||
|
"\n",
|
||||||
|
"print(\"\\nAll forecast jobs have been submitted. Starting real-time status monitor...\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Monitoring loop\n",
|
||||||
|
"completed_jobs = set()\n",
|
||||||
|
"while len(completed_jobs) < len(launched_jobs):\n",
|
||||||
|
" for job in launched_jobs:\n",
|
||||||
|
" if job.name in completed_jobs:\n",
|
||||||
|
" continue\n",
|
||||||
|
"\n",
|
||||||
|
" job._sync_gca_resource()\n",
|
||||||
|
" state = job.state.value if hasattr(job.state, \"value\") else int(job.state)\n",
|
||||||
|
" timestamp = datetime.datetime.now().strftime(\"%H:%M:%S\")\n",
|
||||||
|
" print(\n",
|
||||||
|
" f\"[{timestamp}] Job '{job.display_name}': {job.state.name} (Code: {state})\"\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" # Final states: 4 is SUCCEEDED, 5 is FAILED, 7 is CANCELLED\n",
|
||||||
|
" if state in [4, 5, 7]:\n",
|
||||||
|
" completed_jobs.add(job.name)\n",
|
||||||
|
" if state == 4:\n",
|
||||||
|
" print(f\"\\n🎉 SUCCESS: Job '{job.display_name}' completed successfully!\")\n",
|
||||||
|
" else:\n",
|
||||||
|
" print(\n",
|
||||||
|
" f\"\\n⚠️ FAILURE/CANCEL: Job '{job.display_name}' exited with code {state}. Error: {getattr(job, 'error', 'N/A')}\"\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" if len(completed_jobs) < len(launched_jobs):\n",
|
||||||
|
" time.sleep(20)\n",
|
||||||
|
"\n",
|
||||||
|
"print(\"\\nAll forecast runs are complete!\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "d11c21c4",
|
||||||
|
"metadata": {
|
||||||
|
"id": "112398663144"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"## Visualize Forecasts (Unified) "
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "b04bd730",
|
||||||
|
"metadata": {
|
||||||
|
"id": "x6ZVqWRopWcI"
|
||||||
|
},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# @markdown Select which forecast output you want to visualize. This single component\n",
|
||||||
|
"# @markdown can handle both the standard 6-hourly predictions and the datasets\n",
|
||||||
|
"# @markdown with 1-hour model (which have a 'subtime' dimension).\n",
|
||||||
|
"# @markdown If you run into `Error loading Zarr store: unrecognized engine 'zarr'...` try restarting the runtime session and reruning this cell.\n",
|
||||||
|
"\n",
|
||||||
|
"# @markdown ---\n",
|
||||||
|
"# @markdown ### Visualization Settings\n",
|
||||||
|
"# @markdown - **`model_seed_to_visualize`**: Choose a specific model seed (1-4) to visualize. This should be one of the model seeds selected in the **Forecast Configuration** above.\n",
|
||||||
|
"# @markdown - **`time_steps_to_visualize`**: Choose to visualize 1-hourly or 6-hourly forecasts. If 1-hourly is selected, ensure `enable_hourly_prediction` was selected in the **Forecast Configuration** above.\n",
|
||||||
|
"# @markdown - **`variable_to_visualize`**: Choose the weather variable to visualize. See the [WeatherNext documentation](https://developers.google.com/weathernext/guides/model-specs-vmg) for variable names and descriptions.\n",
|
||||||
|
"# @markdown - **`sample_to_visualize`**: Choose the sample (ensemble member) to visualize.\n",
|
||||||
|
"# @markdown - **`plot_size`**: Choose the size of the plot to generate.\n",
|
||||||
|
"model_seed_to_visualize = \"4\" # @param [\"1\", \"2\", \"3\", \"4\"]\n",
|
||||||
|
"time_steps_to_visualize = \"6-Hourly\" # @param [\"6-Hourly\", \"1-Hourly\"]\n",
|
||||||
|
"variable_to_visualize = \"2m_temperature\" # @param {type:\"string\"}\n",
|
||||||
|
"sample_to_visualize = 0 # @param {type:\"integer\"}\n",
|
||||||
|
"plot_size = 8 # @param {type:\"number\"}\n",
|
||||||
|
"level_to_visualize = None\n",
|
||||||
|
"# @markdown ---\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# --- Helper Functions ---\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def init_time_to_folder_path(init_time: str) -> str:\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" Convert init time to expected GCS folder path.\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" init_date, init_time = init_time.split(\"T\")\n",
|
||||||
|
" return f\"{init_date.replace('-', '')}_{init_time[0:2]}hr\"\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# override these if you'd like to visualize a different set of forecasts\n",
|
||||||
|
"visualize_bucket = BUCKET_URI\n",
|
||||||
|
"visualize_init_date = forecast_init_time\n",
|
||||||
|
"\n",
|
||||||
|
"# set paths based on chosen model seed, bucket, and init date\n",
|
||||||
|
"path_to_6hr_zarr = f\"{BUCKET_URI}/weathernext2_outputs/weathernext_2_seed_{model_seed_to_visualize}/{init_time_to_folder_path(visualize_init_date)}_01_preds/predictions.zarr/\"\n",
|
||||||
|
"path_to_1hr_zarr = f\"{BUCKET_URI}/weathernext2_outputs/weathernext_2_seed_{model_seed_to_visualize}_hourly/{init_time_to_folder_path(visualize_init_date)}_01_preds/predictions.zarr/\"\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"from typing import Optional\n",
|
||||||
|
"\n",
|
||||||
|
"import matplotlib\n",
|
||||||
|
"import matplotlib.animation as animation\n",
|
||||||
|
"import matplotlib.pyplot as plt\n",
|
||||||
|
"import numpy as np\n",
|
||||||
|
"import xarray\n",
|
||||||
|
"from IPython.display import HTML\n",
|
||||||
|
"\n",
|
||||||
|
"matplotlib.rcParams[\"animation.embed_limit\"] = 500\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# --- Helper Functions ---\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def select_data(\n",
|
||||||
|
" data: xarray.Dataset,\n",
|
||||||
|
" variable: str,\n",
|
||||||
|
" level: Optional[int] = None,\n",
|
||||||
|
") -> xarray.Dataset:\n",
|
||||||
|
" \"\"\"Selects a variable from the dataset and optionally a level.\"\"\"\n",
|
||||||
|
" data = data[variable]\n",
|
||||||
|
" if \"batch\" in data.dims:\n",
|
||||||
|
" data = data.isel(batch=0)\n",
|
||||||
|
" if level is not None and \"level\" in data.coords:\n",
|
||||||
|
" data = data.sel(level=level)\n",
|
||||||
|
" return data\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def scale_data(\n",
|
||||||
|
" data: xarray.Dataset,\n",
|
||||||
|
" center: Optional[float] = None,\n",
|
||||||
|
" robust: bool = False,\n",
|
||||||
|
") -> tuple[xarray.Dataset, matplotlib.colors.Normalize, str]:\n",
|
||||||
|
" \"\"\"Scales the data for visualization.\"\"\"\n",
|
||||||
|
" vmin = np.nanpercentile(data.values, (2 if robust else 0))\n",
|
||||||
|
" vmax = np.nanpercentile(data.values, (98 if robust else 100))\n",
|
||||||
|
" if center is not None:\n",
|
||||||
|
" diff = max(vmax - center, center - vmin)\n",
|
||||||
|
" vmin = center - diff\n",
|
||||||
|
" vmax = center + diff\n",
|
||||||
|
" return (\n",
|
||||||
|
" data,\n",
|
||||||
|
" matplotlib.colors.Normalize(vmin, vmax),\n",
|
||||||
|
" (\"RdBu_r\" if center is not None else \"viridis\"),\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def create_forecast_animation(\n",
|
||||||
|
" dataset: xarray.Dataset,\n",
|
||||||
|
" fig_title: str,\n",
|
||||||
|
" plot_size: float = 5,\n",
|
||||||
|
" robust: bool = False,\n",
|
||||||
|
") -> HTML:\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" Creates a forecast animation from an xarray Dataset.\n",
|
||||||
|
" It intelligently handles datasets with or without a 'subtime' dimension.\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" # --- Data Preparation ---\n",
|
||||||
|
" # Check if the data still has 'subtime'). If so, stack dimensions.\n",
|
||||||
|
" # Otherwise, just rename the 'time' dimension for consistency.\n",
|
||||||
|
" if \"subtime\" in dataset.dims:\n",
|
||||||
|
" print(\"Detected 'subtime' dimension. Stacking for hourly animation.\")\n",
|
||||||
|
" # Stack 'time' and 'subtime' into a single animation dimension\n",
|
||||||
|
" plot_data = dataset.stack(animation_step=(\"time\", \"subtime\")).transpose(\n",
|
||||||
|
" \"animation_step\", \"lat\", \"lon\"\n",
|
||||||
|
" )\n",
|
||||||
|
" else:\n",
|
||||||
|
" print(\"No 'subtime' dimension found. Using 'time' for 6-hourly animation.\")\n",
|
||||||
|
" # Use 'time' as the animation dimension\n",
|
||||||
|
" plot_data = dataset.rename({\"time\": \"animation_step\"})\n",
|
||||||
|
"\n",
|
||||||
|
" # Now, the animation dimension is always called 'animation_step'\n",
|
||||||
|
" max_steps = plot_data.sizes[\"animation_step\"]\n",
|
||||||
|
" init_time = plot_data.coords[\"init_time\"].values\n",
|
||||||
|
"\n",
|
||||||
|
" # Scale the data for color mapping\n",
|
||||||
|
" scaled_data, norm, cmap = scale_data(plot_data, robust=robust)\n",
|
||||||
|
"\n",
|
||||||
|
" # --- Plotting Setup ---\n",
|
||||||
|
" figure = plt.figure(figsize=(plot_size * 2, plot_size))\n",
|
||||||
|
" ax = figure.add_subplot(1, 1, 1)\n",
|
||||||
|
" ax.set_xticks([])\n",
|
||||||
|
" ax.set_yticks([])\n",
|
||||||
|
" figure.suptitle(fig_title, fontsize=16)\n",
|
||||||
|
" figure.tight_layout(rect=[0, 0.03, 1, 0.95]) # Adjust for title\n",
|
||||||
|
"\n",
|
||||||
|
" im = ax.imshow(\n",
|
||||||
|
" scaled_data.isel(animation_step=0), norm=norm, origin=\"lower\", cmap=cmap\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" plt.colorbar(\n",
|
||||||
|
" mappable=im,\n",
|
||||||
|
" ax=ax,\n",
|
||||||
|
" orientation=\"vertical\",\n",
|
||||||
|
" pad=0.02,\n",
|
||||||
|
" aspect=16,\n",
|
||||||
|
" shrink=0.75,\n",
|
||||||
|
" cmap=cmap,\n",
|
||||||
|
" extend=(\"both\" if robust else \"neither\"),\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" # --- Animation Update Function ---\n",
|
||||||
|
" def update(frame):\n",
|
||||||
|
" # Get the coordinates for the current frame\n",
|
||||||
|
" step_coords = plot_data[\"animation_step\"][frame].coords\n",
|
||||||
|
"\n",
|
||||||
|
" # Calculate total offset and valid time based on available coordinates\n",
|
||||||
|
" if \"subtime\" in step_coords: # Hourly data\n",
|
||||||
|
" total_offset = step_coords[\"time\"].values + step_coords[\"subtime\"].values\n",
|
||||||
|
" else: # 6-hourly data\n",
|
||||||
|
" total_offset = step_coords[\"animation_step\"].values\n",
|
||||||
|
"\n",
|
||||||
|
" total_hours = total_offset / np.timedelta64(1, \"h\")\n",
|
||||||
|
" valid_time = init_time + total_offset\n",
|
||||||
|
" valid_time_str = np.datetime_as_string(valid_time, unit=\"m\").replace(\"T\", \" \")\n",
|
||||||
|
"\n",
|
||||||
|
" new_title = (\n",
|
||||||
|
" f\"{fig_title}\\n\"\n",
|
||||||
|
" f\"Valid: {valid_time_str} UTC (Forecast: +{total_hours:.1f}h)\"\n",
|
||||||
|
" )\n",
|
||||||
|
" figure.suptitle(new_title, fontsize=16)\n",
|
||||||
|
" im.set_data(scaled_data.isel(animation_step=frame))\n",
|
||||||
|
"\n",
|
||||||
|
" # --- Create and Display Animation ---\n",
|
||||||
|
" ani = animation.FuncAnimation(\n",
|
||||||
|
" fig=figure, func=update, frames=max_steps, interval=250\n",
|
||||||
|
" )\n",
|
||||||
|
" plt.close(figure.number)\n",
|
||||||
|
" return HTML(ani.to_html5_video())\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# --- Main Visualization Logic ---\n",
|
||||||
|
"\n",
|
||||||
|
"# 1. Select the correct path based on the user's dropdown choice\n",
|
||||||
|
"if time_steps_to_visualize == \"6-Hourly\":\n",
|
||||||
|
" path_to_zarr = path_to_6hr_zarr\n",
|
||||||
|
"elif time_steps_to_visualize == \"1-Hourly\":\n",
|
||||||
|
" path_to_zarr = path_to_1hr_zarr\n",
|
||||||
|
"else:\n",
|
||||||
|
" raise ValueError(\"Invalid visualization target selected.\")\n",
|
||||||
|
"\n",
|
||||||
|
"print(f\"Loading data from: {path_to_zarr}\")\n",
|
||||||
|
"\n",
|
||||||
|
"# 2. Load the dataset\n",
|
||||||
|
"try:\n",
|
||||||
|
" full_dataset = xarray.open_zarr(path_to_zarr)\n",
|
||||||
|
"except Exception as e:\n",
|
||||||
|
" print(f\"Error loading Zarr store: {e}\")\n",
|
||||||
|
" # This is a common point of failure, so we exit gracefully.\n",
|
||||||
|
"else:\n",
|
||||||
|
" # 3. Select the specific data slice for visualization\n",
|
||||||
|
" data_for_vis = full_dataset.isel(sample=sample_to_visualize)\n",
|
||||||
|
" variable_data = select_data(data_for_vis, variable_to_visualize, level_to_visualize)\n",
|
||||||
|
"\n",
|
||||||
|
" # 4. Generate the title\n",
|
||||||
|
" title = f\"{variable_to_visualize} (Sample {sample_to_visualize})\"\n",
|
||||||
|
" if level_to_visualize:\n",
|
||||||
|
" title += f\" at {level_to_visualize} hPa\"\n",
|
||||||
|
"\n",
|
||||||
|
" # 5. Create and display the animation\n",
|
||||||
|
" display(create_forecast_animation(variable_data, title, plot_size, robust=True))"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"metadata": {
|
||||||
|
"colab": {
|
||||||
|
"name": "weathernext_2_ic_pc.ipynb",
|
||||||
|
"toc_visible": true
|
||||||
|
},
|
||||||
|
"kernelspec": {
|
||||||
|
"display_name": "Python 3",
|
||||||
|
"name": "python3"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"nbformat": 4,
|
||||||
|
"nbformat_minor": 0
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user