coiled_drift_test.ipynb

coiled_drift_test.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4b6a0c24-dc55-4a0f-a7aa-10e6d1f70c6d",
   "metadata": {},
   "source": [
    "# Run OpenDrift in parallel using Coiled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e6f54e92-04fc-4451-86b7-168d196d7364",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "import coiled\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "\n",
    "import pandas as pd\n",
    "import intake\n",
    "from opendrift.readers import reader_ROMS_native\n",
    "from opendrift.models.oceandrift import OceanDrift\n",
    "import datetime "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1f5e89c7-4a8a-4f3a-a782-55c14f65cde4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1: Get a list of all start_dates\n",
    "granules = pd.date_range(start='2012-01', end='2012-02', freq='28D').to_pydatetime()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0b092106-b00a-4c69-9410-0e123b090202",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([datetime.datetime(2012, 1, 1, 0, 0),\n",
       "       datetime.datetime(2012, 1, 29, 0, 0)], dtype=object)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "granules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "72909cd7-9134-40bb-8bf9-6d9411561cb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 2: Create a function to process each file.\n",
    "#         Load and subset each data granule / file.\n",
    "@coiled.function(\n",
    "    region=\"us-east-1\",                  # Run in the same region as data\n",
    "    spot_policy=\"spot_with_fallback\",    # Use spot instances when available\n",
    "    memory=[\"8GB\", \"16GiB\"],\n",
    "    software='odrift-d6')\n",
    "def process(start_time):\n",
    "\n",
    "    intake_catalog = 'https://renc.osn.xsede.org/rsignellbucket2/rsignell/CNAPS/cnaps_intake.yml'\n",
    "    cat = intake.open_catalog(intake_catalog)\n",
    "    dataset = 'CNAPS_Forecast_Archive_cache'      # all years, read from OSN\n",
    "    dset = cat[dataset]\n",
    "    dset.kwargs['decode_times'] = False\n",
    "    ds = dset.read_chunked()\n",
    "    cnaps = reader_ROMS_native.Reader(ds) \n",
    " \n",
    "    with tempfile.TemporaryDirectory() as tmpdir:\n",
    "        ncfile = f'opendrift_{start_time.strftime(\"%Y%m%d\")}.nc'\n",
    "        o = OceanDrift(loglevel=50)  # Set loglevel to 0 for debug information\n",
    "        o.add_reader(cnaps)\n",
    "        o.seed_elements(lon=-70.462, lat=42.807, radius=1000, number=10, \n",
    "                        z=np.linspace(0, -150, 10), time=start_time)   \n",
    "        o.run(time_step=3600, duration=datetime.timedelta(days=7), outfile=ncfile)\n",
    "    return xr.open_dataset(ncfile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6ce27e14-5801-468c-a09b-5ffa04cf0801",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# run with Distributed locally\n",
    "#from dask.distributed import Client\n",
    "#client = Client(n_workers=2)    # 2 workers @ 8GB on \"large nebari environment\"\n",
    "\n",
    "#futures = client.map(process, granules)\n",
    "#results = client.gather(futures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "be4dbc93-1358-486e-a763-e7d0f9581888",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8602bb5f0c704185a37df708f2adfbd6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Error getting token from client GCP session: Could not automatically determine credentials. Please set GOOGLE_APPLICATION_CREDENTIALS or explicitly create credentials and re-run the application. For more information, please see https://cloud.google.com/docs/authentication/getting-started\n",
      "2024-01-29 08:26:55,408 - distributed.deploy.adaptive - INFO - Adaptive scaling started: minimum=0 maximum=500\n"
     ]
    }
   ],
   "source": [
    "# Step 3: Run Coiled function on each file in parallel\n",
    "results = process.map(granules)     # This runs on the cloud in parallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fcee4615-34f8-497e-a975-34084c6307fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<generator object Function.map.<locals>.<genexpr> at 0x7fd488e2f010>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Step 4: Combine and plot results\n",
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63371e28-cb76-4681-9b60-66b339305741",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "global-global-opendrift",
   "language": "python",
   "name": "conda-env-global-global-opendrift-py"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}