examples/fish_turbulent_flow/Fish_Turbulent_Flow.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "h6tXROWnXnA0",
        "2lPv9zmHafXG"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Fish Navigation in Turbulent Flow"
      ],
      "metadata": {
        "id": "Ps56lybMfTpl"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 321
        },
        "id": "i-BvcB7ufg4m",
        "outputId": "ce0e7fd1-0c85-4b98-de80-f60a053e247c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "<img src=\"https://github.com/Tim-Salzmann/l4casadi/blob/main/examples/fish_turbulent_flow/media/trajectory_generation_vorticity.gif?raw=true\"/>"
            ],
            "text/plain": [
              "<IPython.core.display.Image object>"
            ]
          },
          "metadata": {},
          "execution_count": 1
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install L4CasADi and *Dependencies*"
      ],
      "metadata": {
        "id": "h6tXROWnXnA0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import sys"
      ],
      "metadata": {
        "id": "O5dM1r3WaawT"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "C2LbgtKIWBje",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "09f46c68-98db-46c9-931e-ba2f187745a2"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (2.1.0+cu118)\n",
            "Collecting scikit-build\n",
            "  Downloading scikit_build-0.17.6-py3-none-any.whl (84 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m84.3/84.3 kB\u001b[0m \u001b[31m2.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (3.27.7)\n",
            "Collecting ninja\n",
            "  Downloading ninja-1.11.1.1-py2.py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl (307 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m307.2/307.2 kB\u001b[0m \u001b[31m16.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch) (3.13.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch) (4.5.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch) (3.2.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch) (3.1.2)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch) (2023.6.0)\n",
            "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from torch) (2.1.0)\n",
            "Requirement already satisfied: distro in /usr/lib/python3/dist-packages (from scikit-build) (1.7.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from scikit-build) (23.2)\n",
            "Requirement already satisfied: setuptools>=42.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit-build) (67.7.2)\n",
            "Requirement already satisfied: tomli in /usr/local/lib/python3.10/dist-packages (from scikit-build) (2.0.1)\n",
            "Requirement already satisfied: wheel>=0.32.0 in /usr/local/lib/python3.10/dist-packages (from scikit-build) (0.42.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch) (2.1.3)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch) (1.3.0)\n",
            "Installing collected packages: ninja, scikit-build\n",
            "Successfully installed ninja-1.11.1.1 scikit-build-0.17.6\n",
            "Collecting git+https://github.com/Tim-Salzmann/l4casadi\n",
            "  Cloning https://github.com/Tim-Salzmann/l4casadi to /tmp/pip-req-build-i6mc8eus\n",
            "  Running command git clone --filter=blob:none --quiet https://github.com/Tim-Salzmann/l4casadi /tmp/pip-req-build-i6mc8eus\n",
            "  Resolved https://github.com/Tim-Salzmann/l4casadi to commit 246db23decda71dab14cb83bb462cc1c63e9d7ea\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from l4casadi==1.3.0) (2.1.0+cu118)\n",
            "Collecting casadi>=3.6 (from l4casadi==1.3.0)\n",
            "  Downloading casadi-3.6.4-cp310-none-manylinux2014_x86_64.whl (75.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m75.3/75.3 MB\u001b[0m \u001b[31m9.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: jinja2>=3.1 in /usr/local/lib/python3.10/dist-packages (from l4casadi==1.3.0) (3.1.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from casadi>=3.6->l4casadi==1.3.0) (1.23.5)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2>=3.1->l4casadi==1.3.0) (2.1.3)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch->l4casadi==1.3.0) (3.13.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch->l4casadi==1.3.0) (4.5.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch->l4casadi==1.3.0) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->l4casadi==1.3.0) (3.2.1)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch->l4casadi==1.3.0) (2023.6.0)\n",
            "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from torch->l4casadi==1.3.0) (2.1.0)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch->l4casadi==1.3.0) (1.3.0)\n",
            "Building wheels for collected packages: l4casadi\n",
            "  Building wheel for l4casadi (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for l4casadi: filename=l4casadi-1.3.0-cp310-cp310-linux_x86_64.whl size=46606 sha256=084ce0586a0653793eb26a2ce1ccf7fd90d54ac97b77df795c44b9a31141841c\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-3xit0at0/wheels/cd/54/8a/b8796f827085bb3682fe49796c0f4fe19ddeec1a76ab4187d1\n",
            "Successfully built l4casadi\n",
            "Installing collected packages: casadi, l4casadi\n",
            "Successfully installed casadi-3.6.4 l4casadi-1.3.0\n"
          ]
        }
      ],
      "source": [
        "# @title\n",
        "!pip install torch --index-url https://download.pytorch.org/whl/cpu\n",
        "!pip install scikit-build cmake ninja\n",
        "!pip install git+https://github.com/Tim-Salzmann/l4casadi --no-build-isolation"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title\n",
        "!git clone https://github.com/Tim-Salzmann/l4casadi /tmp/l4casadi\n",
        "sys.path.append('/tmp/l4casadi/examples/fish_turbulent_flow')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EFRHrYHlYHPR",
        "outputId": "10b065b1-c920-4858-f325-206316512c91"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into '/tmp/l4casadi'...\n",
            "remote: Enumerating objects: 437, done.\u001b[K\n",
            "remote: Counting objects: 100% (57/57), done.\u001b[K\n",
            "remote: Compressing objects: 100% (40/40), done.\u001b[K\n",
            "remote: Total 437 (delta 22), reused 41 (delta 15), pack-reused 380\u001b[K\n",
            "Receiving objects: 100% (437/437), 27.07 MiB | 20.43 MiB/s, done.\n",
            "Resolving deltas: 100% (200/200), done.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title\n",
        "from trajectory_generation import trajectory_generator_solver\n",
        "from utils import plot_velocity_field_particle"
      ],
      "metadata": {
        "id": "xkJKdK4QY8k2"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Import"
      ],
      "metadata": {
        "id": "2lPv9zmHafXG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from matplotlib.animation import FuncAnimation\n",
        "import matplotlib.pyplot as plt\n",
        "import casadi as cs\n",
        "import torch\n",
        "import numpy as np"
      ],
      "metadata": {
        "id": "qCHQlRT_ae3z"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import l4casadi as l4c"
      ],
      "metadata": {
        "id": "yW05lzv-aj3H"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Optimization"
      ],
      "metadata": {
        "id": "CIMNrhBdaMxi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Set Fish Start Point Position\n",
        "y_start_pos = 1.5 # @param {type:\"slider\", min:-1.8, max:1.8, step:0.1}\n",
        "p_start = np.array([7.75, y_start_pos])\n",
        "p_goal = np.array([-0.85, -0.4])\n",
        "u_lim = 1\n",
        "T = 20\n",
        "N = 151\n",
        "dt = T / N"
      ],
      "metadata": {
        "id": "CQPMp6farC35",
        "cellView": "form"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Load PyTorch Turbulent Flow Model"
      ],
      "metadata": {
        "id": "K41irO2csaAO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "checkpoint = torch.load(\n",
        "    \"/tmp/l4casadi/examples/fish_turbulent_flow/models/turbolent_flow_model.pt\",\n",
        "    map_location=torch.device('cpu'),\n",
        ")\n",
        "\n",
        "# Standardization\n",
        "model = checkpoint[\"model\"]\n",
        "meanX = checkpoint[\"mean\"][\"x\"]\n",
        "stdX = checkpoint[\"std\"][\"x\"]\n",
        "meanY = checkpoint[\"mean\"][\"y\"]\n",
        "stdY = checkpoint[\"std\"][\"y\"]"
      ],
      "metadata": {
        "id": "i9XOiWOlsNlr"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Create L4CasADi Model from PyTorch Model"
      ],
      "metadata": {
        "id": "-pCgugGrshDT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "x = cs.MX.sym(\"x\", 3)\n",
        "xn = (x - meanX) / stdX\n",
        "\n",
        "y = l4c.L4CasADi(model, name=\"turbulent_model\")(xn)\n",
        "\n",
        "y = y * stdY + meanY\n",
        "fU = cs.Function(\"fU\", [x], [y[0]])\n",
        "fV = cs.Function(\"fV\", [x], [y[1]])"
      ],
      "metadata": {
        "id": "7l7ipOIOslWR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Optimization for Energy Efficiency\n",
        "(This can take 1-2 minutes on Colab CPU)"
      ],
      "metadata": {
        "id": "I11Fh70isqDl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Generate solver\n",
        "nlp = trajectory_generator_solver(\n",
        "    fU=fU, fV=fV, dt=dt, N=N, T=T, u_lim=u_lim, GT=False)\n",
        "\n",
        "# Set Initial Guess and Parameters\n",
        "params = np.vstack([p_start, np.tile(p_goal[:, None], N).T])\n",
        "u_init = np.zeros((N - 1, 2))\n",
        "p_init = np.zeros((N, 2))\n",
        "p_init[:, :] = p_start\n",
        "x_init = np.vstack([p_init, u_init])\n",
        "\n",
        "# Solve NLP\n",
        "x_init_flat = cs.reshape(x_init, 4 * N - 2, 1)\n",
        "params_flat = cs.reshape(params, (N + 1) * 2, 1)\n",
        "sol = nlp[\"solver\"](x0=x_init_flat, p=params_flat, lbg=nlp[\"lbg\"], ubg=nlp[\"ubg\"])\n",
        "\n",
        "# extract solution\n",
        "p_sol = np.squeeze(sol[\"x\"])[: N * 2].reshape(2, N).T\n",
        "u_sol = np.squeeze(sol[\"x\"])[N * 2 :].reshape(2, N - 1).T"
      ],
      "metadata": {
        "id": "-XtEtJDTquCE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Visualize Result"
      ],
      "metadata": {
        "id": "E5m5Gtx2tua6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Generate velocity field for visualization\n",
        "neval = 25\n",
        "Xgrid, Ygrid = np.meshgrid(np.linspace(-1, 8, neval), np.linspace(-2, 2, neval))\n",
        "U = np.zeros((N, neval, neval))\n",
        "V = np.zeros((N, neval, neval))\n",
        "for t in range(0, N):\n",
        "    for i in range(neval):\n",
        "        for j in range(neval):\n",
        "              U[t, i, j] = np.squeeze(fU([t * T / (N - 1), Xgrid[i, j], Ygrid[i, j]]))\n",
        "              V[t, i, j] = np.squeeze(fV([t * T / (N - 1), Xgrid[i, j], Ygrid[i, j]]))"
      ],
      "metadata": {
        "id": "WQHrNHUQu2Ma"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "# Create Animation\n",
        "fig, ax = plt.subplots(figsize=(10, 5.))\n",
        "frames = N\n",
        "anim = FuncAnimation(\n",
        "    fig,\n",
        "    lambda frame_num: plot_velocity_field_particle(\n",
        "        Xgrid,\n",
        "        Ygrid,\n",
        "        U[frame_num],\n",
        "        V[frame_num],\n",
        "        p_sol[max(0, frame_num - 10): frame_num + 1, 0],\n",
        "        p_sol[max(0, frame_num - 10): frame_num + 1, 1],\n",
        "        p_start,\n",
        "        p_goal,\n",
        "        round(frame_num / frames * T, 3),\n",
        "    ),\n",
        "    frames=frames,\n",
        "    interval=100,\n",
        ")\n",
        "\n",
        "anim.save('anim.gif')"
      ],
      "metadata": {
        "id": "UHLP7CKutt6k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from IPython.display import Image\n",
        "Image('./anim.gif')"
      ],
      "metadata": {
        "id": "y-yE8qox47df"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "iEpT51Au4riq"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}