eddy_identification_paralelised_share.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Eddy ID and tracking on all avalibale 2D ICON EERIE data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The aim of this notebook is to run parallelised eddy ID and tracking on model data on levante. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook uses the \"updated\" version of py-eedy-tracker that was improved by Aaron Wienkers during the first eerie hackathon (Nov 2023). Link to the py-eddy-tracker updated version and how to use it - https://github.com/eerie-project/EERIE_hackathon_2023/tree/main/RESULTS/pyeddytracker_xarray_dask_parallel. The updated version is able to use intake catalongs and xarrays.\n",
    "\n",
    "\n",
    "Most of the Eddy ID and tracking functions come from py-eddy-tracker (https://py-eddy-tracker.readthedocs.io/en/stable/) and a lot of other eddy ID examples from EERIE first hackathon (Nov 2023) (https://github.com/eerie-project/EERIE_hackathon_2023/tree/main/RESULTS). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Regridded 0.25deg ICON 1950 run used (because it is regridded, any model should work). Tracks around 2 million eddies in 39 years of data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Date created: 12.12.2023\n",
    "\n",
    "Last edited: 22.05.2024\n",
    "\n",
    "Stella Bērziņa (stella.berzina@usys.ethz.ch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "import matplotlib.pylab as plt\n",
    "import matplotlib.cm as cm\n",
    "from scipy.interpolate import CloughTocher2DInterpolator, LinearNDInterpolator, NearestNDInterpolator\n",
    "import glob\n",
    "import intake\n",
    "from pathlib import Path\n",
    "import dask\n",
    "import pandas as pd\n",
    "dask.config.set({\"array.slicing.split_large_chunks\": True}) \n",
    "# import cartopy.crs as ccrs\n",
    "# import cartopy.feature as cfeature\n",
    "import sys\n",
    "from numpy import arange, ones\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.lines as mlines\n",
    "from matplotlib.cm import get_cmap\n",
    "from matplotlib import colors \n",
    "from matplotlib.colors import ListedColormap\n",
    "from matplotlib.figure import Figure\n",
    "\n",
    "from py_eddy_tracker.observations.network import NetworkObservations\n",
    "from py_eddy_tracker.observations.observation import EddiesObservations, Table\n",
    "from py_eddy_tracker.observations.tracking import TrackEddiesObservations\n",
    "from py_eddy_tracker.dataset.grid import RegularGridDataset\n",
    "from py_eddy_tracker.featured_tracking.area_tracker import AreaTracker\n",
    "from py_eddy_tracker.gui import GUI\n",
    "from py_eddy_tracker.tracking import Correspondances\n",
    "\n",
    "\n",
    "from datetime import datetime, timedelta\n",
    "from netCDF4 import Dataset\n",
    "\n",
    "import io\n",
    "import os\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "# from concurrent.futures import ProcessPoolExecutor\n",
    "# from concurrent.futures import ThreadPoolExecutor\n",
    "\n",
    "import dask\n",
    "from dask_jobqueue import SLURMCluster\n",
    "from dask.distributed import Client\n",
    "\n",
    "import cartopy.crs as ccrs\n",
    "from cartopy.mpl.ticker import (LongitudeFormatter, LatitudeFormatter,LatitudeLocator)\n",
    "import cartopy.feature as cfeature\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "grid = xr.open_dataset('/pool/data/ICON/grids/public/mpim/0016/icon_grid_0016_R02B09_O.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get model data to ID eddies from"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['2d_daily_mean', '2d_daily_mean_vertical-remap025', '2d_daily_square', '2d_monthly_mean', '2d_monthly_mean_vertical-remap025', '2d_monthly_square', '5lev_daily_mean', 'eddy_monthly_mean']\n"
     ]
    }
   ],
   "source": [
    "cat = intake.open_catalog(\"https://raw.githubusercontent.com/eerie-project/intake_catalogues/main/eerie.yaml\")\n",
    "expid = 'eerie-control-1950'\n",
    "gridspec = 'gr025'\n",
    "cat_regrid = cat['dkrz.disk.model-output.icon-esm-er'][expid]['ocean'][gridspec]\n",
    "print(list(cat_regrid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds = cat_regrid['2d_daily_mean'].to_dask()\n",
    "varname = 'ssh'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nrelevant if you want to run analysis on a subset of data. we are using all the time teps here\\n'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Relevant if you want to run analysis on a subset of data. I am using all the time steps here.\n",
    "\"\"\"\n",
    "\n",
    "# # 2022-2032\n",
    "# ds_subset = ds.sel(time=slice('2032-01-02','2039-12-31'))\n",
    "# datearr = [pd.Timestamp(t).to_pydatetime() for t in ds_subset.time.values]\n",
    "# ds_subset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Eddy ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/scratch/b/b382618/try/eerie-control-1950/gr025/\n"
     ]
    }
   ],
   "source": [
    "# Directories\n",
    "scratch = '/scratch/b/b382618/try/'\n",
    "\n",
    "datadir = scratch+expid+'/'+gridspec+'/'\n",
    "print(datadir)\n",
    "\n",
    "eddy_dir = datadir+'/eddyID/'\n",
    "if not os.path.exists(eddy_dir):\n",
    "    os.makedirs(eddy_dir)\n",
    "\n",
    "track_dir = datadir+'/eddyTracks/'\n",
    "if not os.path.exists(track_dir):\n",
    "    os.makedirs(track_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Start SLURMCluster on Levante for Parallel Computing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define cluster; corresponds to one SLURM job\n",
    "cluster = SLURMCluster(name='dask-cluster',\n",
    "                      cores=16,  \n",
    "                      memory='256GB',\n",
    "                      processes=16,\n",
    "                      interface='ib0',\n",
    "                      queue='compute',\n",
    "                      account='bk1377',\n",
    "                      walltime='05:00:00',\n",
    "                      asynchronous=0,\n",
    "                      log_directory='/home/b/b382618/.log_trash')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-5a3e28be-cdaf-11ee-89a3-080038c03255</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> dask_jobqueue.SLURMCluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"http://136.172.120.76:37223/status\" target=\"_blank\">http://136.172.120.76:37223/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">SLURMCluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">dask-cluster</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"http://136.172.120.76:37223/status\" target=\"_blank\">http://136.172.120.76:37223/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 0\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 0\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 0 B\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-3bcc3ee3-ef0f-425f-b3c7-e4d527621acb</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tcp://136.172.120.76:43205\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://136.172.120.76:37223/status\" target=\"_blank\">http://136.172.120.76:37223/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 0 B\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tcp://136.172.120.76:43205' processes=0 threads=0, memory=0 B>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# multiply cluster size by 2; equivalent to 2 SLURM jobs\n",
    "# this will actually open 2 jobs on squeue\n",
    "\n",
    "cluster.scale(cores=32)\n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Eddy ID parameters used by AVISO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "wavelength = 700 #choice of spatial cutoff for high pass filter in km\n",
    "step_ht = 0.002 #intervals to search for closed contours (2mm in this case)\n",
    "pixel_limit = None  # Min and max pixel count for valid contour\n",
    "shape_error = 70  # Error max (%) between ratio of circle fit and contour"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def detection(ncfile, varname, date): \n",
    "    g = RegularGridDataset(None, \"lon\", \"lat\", centered=True, nc4file=ncfile)  # NOTE: Using 'None' for the .nc file path then requires specifying directly the netcdf4 variable in memory\n",
    "    g.add_uv(varname)\n",
    "    g.bessel_high_filter(varname, wavelength, order=1)\n",
    "        \n",
    "    a, c = g.eddy_identification(varname, \"u\", \"v\", \n",
    "    date,  # Date of identification\n",
    "    step_ht,  # step between two isolines of detection (m)\n",
    "    pixel_limit=pixel_limit,  # Min and max pixel count for valid contour\n",
    "    shape_error=70  # Error max (%) between ratio of circle fit and contour\n",
    "    )\n",
    "    return a,c,g\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def detection_save_netcdf_output(date, tt):\n",
    "    # Load data from xarray into netcdf4 type\n",
    "    da_ssh = ds_subset.isel(time=tt).ssh.drop('time')\n",
    "    da_ssh.load(scheduler='sync')\n",
    "    da_netcdf = Dataset('in-mem-file', mode='r', memory=da_ssh.to_netcdf())\n",
    "    \n",
    "    print('Identifying daily eddies for '+date.strftime('%Y%m%d'))\n",
    "    a_filtered, c_filtered, g_filtered = detection(da_netcdf,varname,date)\n",
    "    with Dataset(date.strftime(eddy_dir+\"eddyID_anticyclonic_\"+date.strftime('%Y%m%d')+\".nc\"), \"w\") as h:\n",
    "        a_filtered.to_netcdf(h)\n",
    "    with Dataset(date.strftime(eddy_dir+\"eddyID_cyclonic_\"+date.strftime('%Y%m%d')+\".nc\"), \"w\") as h:\n",
    "        c_filtered.to_netcdf(h)\n",
    "    del a_filtered\n",
    "    del c_filtered\n",
    "    del g_filtered\n",
    "    del date"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "lazy_results = []\n",
    "for i in range(len(datearr)):\n",
    "    lazy_result = dask.delayed(detection_save_netcdf_output)(date=datearr[i], tt=int(i))\n",
    "    lazy_results.append(lazy_result)  \n",
    "    \n",
    "futures = dask.compute(*lazy_results)\n",
    "results = dask.compute(*futures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()\n",
    "client.shutdown()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Eddy tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from py_eddy_tracker.featured_tracking.area_tracker import AreaTracker\n",
    "from py_eddy_tracker.tracking import Correspondances"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Eddy tracking parameters used by AVISO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "nb_obs_min = 10 # minimum of 10 points in track to be considered a long trajectory\n",
    "raw = False # \n",
    "cmin = 0.05 # minimum contour\n",
    "virtual = 4 # number of consecutive timesteps with missing detection allowed\n",
    "class_kw = dict(cmin=cmin)\n",
    "zarr = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# if the number of cores is increased then not enough memory per worker to deal with 37 years of data\n",
    "\n",
    "cluster = SLURMCluster(name='dask-cluster',\n",
    "                      cores=4,\n",
    "                      memory='512GB',\n",
    "                      processes=4,\n",
    "                      interface='ib0',\n",
    "                      queue='compute',\n",
    "                      account='bk1377',\n",
    "                      walltime='06:00:00',\n",
    "                      asynchronous=0,\n",
    "                      log_directory='/home/b/b382618/.log_trash')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-a343cfad-cfea-11ee-ab2e-080038c03229</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> dask_jobqueue.SLURMCluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"http://136.172.120.78:8787/status\" target=\"_blank\">http://136.172.120.78:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">SLURMCluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">dask-cluster</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"http://136.172.120.78:8787/status\" target=\"_blank\">http://136.172.120.78:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 0\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 0\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 0 B\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-1e4df590-6453-422b-90fa-ff8e9d9da922</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tcp://136.172.120.78:37905\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://136.172.120.78:8787/status\" target=\"_blank\">http://136.172.120.78:8787/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 0 B\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tcp://136.172.120.78:37905' processes=0 threads=0, memory=0 B>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cluster.scale(cores=4)\n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Functions from eddy-tracking.py (aided by Malcolm Roberts)\n",
    "\n",
    "def tracking(file_objects, previous_correspondance, eddy_type, zarr=False, nb_obs_min=10, raw=True, cmin=0.05, virtual=4):\n",
    "    # We run a tracking with a tracker which uses contour overlap, on first time step\n",
    "    output_dir = os.path.dirname(previous_correspondance)\n",
    "    class_kw = dict(cmin=cmin)\n",
    "    if not os.path.isfile(previous_correspondance):\n",
    "        c = Correspondances(\n",
    "            datasets=file_objects, class_method=AreaTracker, \n",
    "            class_kw=class_kw, virtual=virtual\n",
    "        )\n",
    "        c.track()\n",
    "        c.prepare_merging()\n",
    "    else:\n",
    "        c = Correspondances(\n",
    "            datasets=file_objects, class_method=AreaTracker, \n",
    "            class_kw=class_kw, virtual=virtual,\n",
    "            previous_correspondance=previous_correspondance\n",
    "        )\n",
    "        c.track()\n",
    "        c.prepare_merging()\n",
    "        c.merge()\n",
    "\n",
    "    new_correspondance = previous_correspondance[:-3]+'_new.nc'\n",
    "    with Dataset(new_correspondance, \"w\") as h:\n",
    "        c.to_netcdf(h)\n",
    "\n",
    "    try:\n",
    "        # test can read new file, and then move to replace old file\n",
    "        nc = Dataset(new_correspondance, 'r')\n",
    "        os.rename(new_correspondance, previous_correspondance)\n",
    "    except:\n",
    "        raise Exception('Error opening new correspondance file '+new_correspondance)\n",
    "\n",
    "    write_obs_files(c, raw, output_dir, zarr, eddy_type, nb_obs_min)\n",
    "    \n",
    "\n",
    "def write_obs_files(c, raw, output_dir, zarr, eddy_type, nb_obs_min):\n",
    "    kw_write = dict(path=output_dir, zarr_flag=zarr, sign_type=eddy_type)\n",
    "\n",
    "    fout = os.path.join(output_dir, eddy_type+'_untracked.nc')\n",
    "    c.get_unused_data(raw_data=raw).write_file(\n",
    "        filename=fout\n",
    "    )\n",
    "\n",
    "    short_c = c._copy()\n",
    "    short_c.shorter_than(size_max=nb_obs_min)\n",
    "    short_track = short_c.merge(raw_data=raw)\n",
    "\n",
    "    if c.longer_than(size_min=nb_obs_min) is False:\n",
    "        long_track = short_track.empty_dataset()\n",
    "    else:\n",
    "        long_track = c.merge(raw_data=raw)\n",
    "\n",
    "    # We flag obs\n",
    "    if c.virtual:\n",
    "        long_track[\"virtual\"][:] = long_track[\"time\"] == 0\n",
    "        long_track.normalize_longitude()\n",
    "        long_track.filled_by_interpolation(long_track[\"virtual\"] == 1)\n",
    "        short_track[\"virtual\"][:] = short_track[\"time\"] == 0\n",
    "        short_track.normalize_longitude()\n",
    "        short_track.filled_by_interpolation(short_track[\"virtual\"] == 1)\n",
    "\n",
    "    print(\"Longer track saved have %d obs\", c.nb_obs_by_tracks.max())\n",
    "    print(\n",
    "        \"The mean length is %d observations for long track\",\n",
    "        c.nb_obs_by_tracks.mean(),\n",
    "    )\n",
    "\n",
    "    fout = os.path.join(output_dir, eddy_type+'_tracks.nc')\n",
    "    long_track.write_file(filename=fout)\n",
    "    fout = os.path.join(output_dir, eddy_type+'_short.nc')\n",
    "    short_track.write_file(\n",
    "        #filename=\"%(path)s/%(sign_type)s_track_too_short.nc\", **kw_write\n",
    "        filename=fout\n",
    "    )\n",
    "\n",
    "\n",
    "\n",
    "def tracking_eddytype(eddy_type, track_dir, eddy_dir, zarr, nb_obs_min, raw, cmin):\n",
    "    previous_correspondance = os.path.join(track_dir, eddy_type+'_correspondance.nc')\n",
    "    search = os.path.join(eddy_dir+'eddyID_'+eddy_type+'_'+'????????.nc')\n",
    "    print('search files ',search)\n",
    "    file_objects = sorted(glob.glob(search))\n",
    "    tracking(file_objects, previous_correspondance, eddy_type, zarr=zarr, nb_obs_min=nb_obs_min, raw=raw, cmin=cmin)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "eddy_dir =  /scratch/b/b382618/try/eerie-control-1950/gr025//eddyID/\n",
      "track_dir =  /scratch/b/b382618/try/eerie-control-1950/gr025//eddyTracks/\n"
     ]
    }
   ],
   "source": [
    "lazy_results = []\n",
    "print('eddy_dir = ', eddy_dir)\n",
    "print('track_dir = ', track_dir)\n",
    "for eddy_type in ['cyclonic','anticyclonic']:\n",
    "    lazy_result = dask.delayed(tracking_eddytype)(track_dir=track_dir,\n",
    "                                                    eddy_dir=eddy_dir,\n",
    "                                                    eddy_type=eddy_type,\n",
    "                                                    zarr=zarr, nb_obs_min=nb_obs_min,\n",
    "                                                    raw=raw, cmin=cmin)\n",
    "    lazy_results.append(lazy_result)  \n",
    "\n",
    "futures = dask.compute(*lazy_results)\n",
    "results = dask.compute(*futures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "client.close()\n",
    "client.shutdown()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Subsetting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I do this subsetting because in a lot of analysis you don't need all tracked eddies but only the long-lived ones. and the dataset is significantly smaller then."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I choose to use the py-eddy-tracker file format because it is pretty useful if you want to create some of the plots that the package provides code for. Esentially "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Tracked eddies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# this file contains all tracked anticyclonic eddies for 37 years\n",
    "a = TrackEddiesObservations.load_file((\n",
    "        \"/work/bk1377/b382618/icon_malcom/10d/eerie-control-1950/gr025/eddyTracks/anticyclonic_tracks.nc\")\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "c = TrackEddiesObservations.load_file(\n",
    "    (\"/work/bk1377/b382618/icon_malcom/10d/eerie-control-1950/gr025/eddyTracks/cyclonic_tracks.nc\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    | 55045328 observations from 18993.999305555553 to 32871.99930555555 (13879.0 days, ~3966 obs/day)\n",
      "    |   Speed area      : 36.36 Mkm²/day\n",
      "    |   Effective area  : 48.69 Mkm²/day\n",
      "    ----Distribution in Amplitude:\n",
      "    |   Amplitude bounds (cm)        0.00      1.00      2.00      3.00      4.00      5.00     10.00    500.00\n",
      "    |   Percent of eddies         :      29.66     26.30     13.78      8.21      5.22     10.58      6.26\n",
      "    ----Distribution in Radius:\n",
      "    |   Speed radius (km)            0.00     15.00     30.00     45.00     60.00     75.00    100.00    200.00   2000.00\n",
      "    |   Percent of eddies         :       1.22     22.49     29.97     18.86     11.41      9.71      6.20      0.13\n",
      "    |   Effective radius (km)        0.00     15.00     30.00     45.00     60.00     75.00    100.00    200.00   2000.00\n",
      "    |   Percent of eddies         :       0.96     17.11     25.34     18.71     12.98     13.19     11.38      0.32\n",
      "    ----Distribution in Latitude\n",
      "        Latitude bounds            -90.00    -60.00    -15.00     15.00     60.00     90.00\n",
      "        Percent of eddies         :      11.61     42.45     10.27     26.34      9.34\n",
      "        Percent of speed area     :       4.21     43.73     25.21     24.58      2.27\n",
      "        Percent of effective area :       4.30     44.66     24.16     24.61      2.26\n",
      "        Mean speed radius (km)    :      32.09     52.92     83.06     50.16     26.06\n",
      "        Mean effective radius (km):      37.42     61.98     95.43     58.31     29.96\n",
      "        Mean amplitude (cm)       :       2.13      4.31      1.79      2.99      2.08\n",
      "    | 1013796 tracks (54.30 obs/tracks, shorter 10 obs, longer 3024 obs)\n",
      "    |   5187637 filled observations (5.12 obs/tracks, 9.42 % of total)\n",
      "    |   Intepolated speed area      : 2.73 Mkm²/day\n",
      "    |   Intepolated effective area  : 3.24 Mkm²/day\n",
      "    |   Distance by day             : Mean 8.14 , Median 4.16 km/day\n",
      "    |   Distance by track           : Mean 433.90 , Median 206.53 km/track\n",
      "    ----Distribution in lifetime:\n",
      "    |   Lifetime (days  )            1.00     30.00     90.00    180.00    270.00    365.00   1000.00  10000.00\n",
      "    |   Percent of tracks         :      51.01     33.00     10.77      3.13      1.18      0.88      0.03\n",
      "    |   Percent of eddies         :      16.22     31.11     24.66     12.50      6.72      8.06      0.73\n"
     ]
    }
   ],
   "source": [
    "print(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nyou can also open the tracked files as an xarray\\n'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "you can also open the tracked files as an xarray\n",
    "\"\"\"\n",
    "\n",
    "# a_xr = xr.open_dataset('/scratch/b/b382618/try/eerie-control-1950/gr025/eddyTracks/anticyclonic_tracks.nc')\n",
    "# c_xr = xr.open_dataset('/scratch/b/b382618/try/eerie-control-1950/gr025/eddyTracks/cyclonic_tracks.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### I downloaded aviso observations from the eerie intake catalog do I could open the output as a py-eddy-tracker file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_obs = TrackEddiesObservations.load_file((\n",
    "        \"/work/bk1377/b382618/icon_malcom/10d/eerie-control-1950/gr025/aviso/anticyclonic.nc\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "c_obs = TrackEddiesObservations.load_file((\n",
    "        \"/work/bk1377/b382618/icon_malcom/10d/eerie-control-1950/gr025/aviso/cyclonic.nc\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    | 34521490 observations from 15706.0 to 26337.0 (10632.0 days, ~3247 obs/day)\n",
      "    |   Speed area      : 32.95 Mkm²/day\n",
      "    |   Effective area  : 45.02 Mkm²/day\n",
      "    ----Distribution in Amplitude:\n",
      "    |   Amplitude bounds (cm)        0.00      1.00      2.00      3.00      4.00      5.00     10.00    500.00\n",
      "    |   Percent of eddies         :      14.24     23.74     15.91     10.63      7.42     17.34     10.73\n",
      "    ----Distribution in Radius:\n",
      "    |   Speed radius (km)            0.00     15.00     30.00     45.00     60.00     75.00    100.00    200.00   2000.00\n",
      "    |   Percent of eddies         :       0.02      8.09     35.82     24.60     13.22     10.72      7.38      0.14\n",
      "    |   Effective radius (km)        0.00     15.00     30.00     45.00     60.00     75.00    100.00    200.00   2000.00\n",
      "    |   Percent of eddies         :       0.02      5.66     27.03     23.13     15.07     14.92     13.75      0.43\n",
      "    ----Distribution in Latitude\n",
      "        Latitude bounds            -90.00    -60.00    -15.00     15.00     60.00     90.00\n",
      "        Percent of eddies         :       6.62     49.64     10.94     28.98      3.83\n",
      "        Percent of speed area     :       3.11     45.23     23.51     26.99      1.15\n",
      "        Percent of effective area :       3.21     46.12     22.27     27.26      1.14\n",
      "        Mean speed radius (km)    :      40.06     54.23     84.19     54.58     32.96\n",
      "        Mean effective radius (km):      47.13     63.74     96.03     63.88     37.83\n",
      "        Mean amplitude (cm)       :       3.60      5.43      2.14      4.74      2.87\n",
      "    | 805895 tracks (42.84 obs/tracks, shorter 10 obs, longer 2447 obs)\n",
      "    |   2517734 filled observations (3.12 obs/tracks, 7.29 % of total)\n",
      "    |   Intepolated speed area      : 2.26 Mkm²/day\n",
      "    |   Intepolated effective area  : 2.68 Mkm²/day\n",
      "    |   Distance by day             : Mean 6.83 , Median 4.33 km/day\n",
      "    |   Distance by track           : Mean 285.55 , Median 165.88 km/track\n",
      "    ----Distribution in lifetime:\n",
      "    |   Lifetime (days  )            1.00     30.00     90.00    180.00    270.00    365.00   1000.00  10000.00\n",
      "    |   Percent of tracks         :      57.95     31.57      7.77      1.74      0.57      0.39      0.02\n",
      "    |   Percent of eddies         :      23.17     36.84     22.14      8.75      4.08      4.51      0.50\n"
     ]
    }
   ],
   "source": [
    "print(a_obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Select only the eddies that have a lifetime >= 16 weeks (so you can compare to Chelton et al. 2016)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_16 = a.extract_with_length((7 * 16, -1))\n",
    "c_16 = c.extract_with_length((7 * 16, -1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "a_obs16 = a_obs.extract_with_length((7 * 16, -1))\n",
    "c_obs16 = c_obs.extract_with_length((7 * 16, -1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "py_eddy_tracker.observations.tracking.TrackEddiesObservations"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(a_16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save to a directory of your choice. to_netcdf does not really work although it is a function in py-eddy-tracker."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "a_16.write_file(\"/work/bk1377/b382618/icon_malcom/16w/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "c_16.write_file(\"/work/bk1377/b382618/icon_malcom/16w/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'a_obs16' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43ma_obs16\u001b[49m\u001b[38;5;241m.\u001b[39mwrite_file(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/work/bk1377/b382618/icon_malcom/16w/aviso/\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'a_obs16' is not defined"
     ]
    }
   ],
   "source": [
    "a_obs16.write_file(\"/work/bk1377/b382618/icon_malcom/16w/aviso/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "c_obs16.write_file(\"/work/bk1377/b382618/icon_malcom/16w/aviso/\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py39.2",
   "language": "python",
   "name": "py39"
  },
  "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}