{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Network segmentation process\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# sphinx_gallery_thumbnail_number = 2\nimport re\n\nfrom matplotlib import pyplot as plt\nfrom matplotlib.animation import FuncAnimation\nfrom matplotlib.colors import ListedColormap\nfrom numpy import ones, where\n\nfrom py_eddy_tracker.data import get_demo_path\nfrom py_eddy_tracker.gui import GUI_AXES\nfrom py_eddy_tracker.observations.network import NetworkObservations\nfrom py_eddy_tracker.observations.tracking import TrackEddiesObservations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "class VideoAnimation(FuncAnimation):\n    def _repr_html_(self, *args, **kwargs):\n        \"\"\"To get video in html and have a player\"\"\"\n        content = self.to_html5_video()\n        return re.sub(\n            r'width=\"[0-9]*\"\\sheight=\"[0-9]*\"', 'width=\"100%\" height=\"100%\"', content\n        )\n\n    def save(self, *args, **kwargs):\n        if args[0].endswith(\"gif\"):\n            # In this case gif is used to create thumbnail which is not used but consume same time than video\n            # So we create an empty file, to save time\n            with open(args[0], \"w\") as _:\n                pass\n            return\n        return super().save(*args, **kwargs)\n\n\ndef get_obs(dataset):\n    \"Function to isolate a specific obs\"\n    return where(\n        (dataset.lat > 33)\n        * (dataset.lat < 34)\n        * (dataset.lon > 22)\n        * (dataset.lon < 23)\n        * (dataset.time > 20630)\n        * (dataset.time < 20650)\n    )[0][0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Hack to pick up each step of segmentation\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "TRACKS = list()\nINDICES = list()\n\n\nclass MyTrack(TrackEddiesObservations):\n    @staticmethod\n    def get_next_obs(i_current, ids, x, y, time_s, time_e, time_ref, window, **kwargs):\n        TRACKS.append(ids[\"track\"].copy())\n        INDICES.append(i_current)\n        return TrackEddiesObservations.get_next_obs(\n            i_current, ids, x, y, time_s, time_e, time_ref, window, **kwargs\n        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Load data\nLoad data where observations are put in same network but no segmentation\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Get a known network for the demonstration\nn = NetworkObservations.load_file(get_demo_path(\"network_med.nc\")).network(651)\n# We keep only some segment\nn = n.relative(get_obs(n), order=2)\nprint(len(n))\n# We convert and order object like segmentation was never happen on observations\ne = n.astype(MyTrack)\ne.obs.sort(order=(\"track\", \"time\"), kind=\"stable\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Do segmentation\nSegmentation based on maximum overlap, temporal window for candidates = 5 days\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "matrix = e.split_network(intern=False, window=5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Anim\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def update(i_frame):\n    tr = TRACKS[i_frame]\n    mappable_tracks.set_array(tr)\n    s = 40 * ones(tr.shape)\n    s[tr == 0] = 4\n    mappable_tracks.set_sizes(s)\n\n    indices_frames = INDICES[i_frame]\n    mappable_CONTOUR.set_data(\n        e.contour_lon_e[indices_frames], e.contour_lat_e[indices_frames]\n    )\n    mappable_CONTOUR.set_color(cmap.colors[tr[indices_frames] % len(cmap.colors)])\n    return (mappable_tracks,)\n\n\nfig = plt.figure(figsize=(16, 9), dpi=60)\nax = fig.add_axes([0.04, 0.06, 0.94, 0.88], projection=GUI_AXES)\nax.set_title(f\"{len(e)} observations to segment\")\nax.set_xlim(19, 29), ax.set_ylim(31, 35.5), ax.grid()\nvmax = TRACKS[-1].max()\ncmap = ListedColormap([\"gray\", *e.COLORS[:-1]], name=\"from_list\", N=vmax)\nmappable_tracks = ax.scatter(\n    e.lon, e.lat, c=TRACKS[0], cmap=cmap, vmin=0, vmax=vmax, s=20\n)\nmappable_CONTOUR = ax.plot(\n    e.contour_lon_e[INDICES[0]], e.contour_lat_e[INDICES[0]], color=cmap.colors[0]\n)[0]\nani = VideoAnimation(fig, update, frames=range(1, len(TRACKS), 4), interval=125)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Final Result\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig = plt.figure(figsize=(16, 9))\nax = fig.add_axes([0.04, 0.06, 0.94, 0.88], projection=GUI_AXES)\nax.set_xlim(19, 29), ax.set_ylim(31, 35.5), ax.grid()\n_ = ax.scatter(e.lon, e.lat, c=TRACKS[-1], cmap=cmap, vmin=0, vmax=vmax, s=20)"
      ]
    }
  ],
  "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.10.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}