{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Graviton SDK: Fluid Dynamics & Weather Research and Forecasting (WRF)\n",
    "This notebook demonstrates how to evaluate WRF 4.8.0 atmospheric modeling workloads, upload a NetCDF virtual world file using `gv.load_WRF()`, and execute multi-node MPI simulations on Google Cloud Batch.\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/CodeZeroLabs/graviton-examples/blob/master/graviton_sdk_wrf_example.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Install the Graviton SDK\n",
    "!pip install graviton-sdk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import graviton as gv\n",
    "import tempfile\n",
    "import os\n",
    "\n",
    "# 2. Authenticate with Graviton\n",
    "gv.login(\"a-little-more-art-than-science\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Defining the WRF Simulation Domain\n",
    "WRF simulation payloads require grid resolution parameters (`dx`, `dy`), grid dimensions (`e_we`, `e_sn`), and simulation duration (`sim_time`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wrf_payload = {\n",
    "    \"t\": \"wrf\",\n",
    "    \"dx\": 10000,      # Grid spacing in X direction (10 km)\n",
    "    \"dy\": 10000,      # Grid spacing in Y direction (10 km)\n",
    "    \"e_we\": 100,      # Grid dimension West-East\n",
    "    \"e_sn\": 100,      # Grid dimension South-North\n",
    "    \"sim_time\": 86400 # Simulation duration in seconds (24 hours)\n",
    "}\n",
    "\n",
    "# Evaluate domain complexity and obtain signed upload URL for the world file\n",
    "eval_data = gv.evaluate(wrf_payload)\n",
    "print(f\"Job ID: {eval_data.id}\")\n",
    "print(f\"Estimated Credit Cost: {eval_data.cost:.2f} Credits\")\n",
    "print(f\"Estimated Duration: {eval_data.estimated_time:.2f}s\")\n",
    "print(f\"Signed Upload URL: {eval_data.upload_url[:50]}...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Uploading the NetCDF Virtual World File\n",
    "WRF atmospheric simulations run inside a virtual world defined in a NetCDF binary file (`world.nc`). The SDK provides `gv.load_WRF(job_id, world_file_path)` which securely streams the NetCDF file directly to Graviton storage via the evaluation's signed URL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a sample NetCDF world file for demonstration\n",
    "temp_world_file = os.path.join(tempfile.gettempdir(), \"demo_world.nc\")\n",
    "with open(temp_world_file, \"wb\") as f:\n",
    "    f.write(b\"CDF\\x02\" + b\"\\x00\" * 1024)\n",
    "\n",
    "# Upload the world file to the job's signed URL\n",
    "print(f\"Uploading world file '{temp_world_file}' for Job '{eval_data.id}'...\")\n",
    "upload_status = gv.load_WRF(eval_data.id, temp_world_file)\n",
    "print(f\"Upload Status: {upload_status}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Launching the Multi-Node MPI Supercomputing Job\n",
    "Once the world file has been uploaded, initiate compute on Google Cloud Batch across MPI-interconnected C2 nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Trigger MPI cluster compute\n",
    "job = gv.compute(eval_data.id)\n",
    "print(f\"Job Status: {job.status()}\")\n",
    "\n",
    "# Wait for completion and view results\n",
    "result = job.wait()\n",
    "print(\"Simulation Completed!\")\n",
    "print(f\"Output URL: {result.get('output_url')}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}