sim.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f9fe8b90",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "from jax import grad, jit, vmap, jacfwd\n",
    "import dhutils.dhutils as dhu\n",
    "import pythreejs as THREE\n",
    "import matplotlib.pyplot as plt\n",
    "import sys\n",
    "sys.path.append(\"./\")\n",
    "from solvers.mpcg import MPCGSolver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4c1e3d3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Implementation based on ideas from\n",
    "# https://cg.informatik.uni-freiburg.de/course_notes/sim_03_cloth1.pdf\n",
    "\n",
    "# define number of edges along each direction\n",
    "# assumption: square cloth\n",
    "N = 4\n",
    "\n",
    "# generate the mesh and get (u, v), (x, y, z) coordinates\n",
    "Nx = N\n",
    "Ny = N"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d428c371",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "faces:\n",
      "[[ 0  1  6]\n",
      " [ 0  6  5]\n",
      " [ 1  2  7]\n",
      " [ 1  7  6]\n",
      " [ 2  3  8]\n",
      " [ 2  8  7]\n",
      " [ 3  4  9]\n",
      " [ 3  9  8]\n",
      " [ 5  6 11]\n",
      " [ 5 11 10]\n",
      " [ 6  7 12]\n",
      " [ 6 12 11]\n",
      " [ 7  8 13]\n",
      " [ 7 13 12]\n",
      " [ 8  9 14]\n",
      " [ 8 14 13]\n",
      " [10 11 16]\n",
      " [10 16 15]\n",
      " [11 12 17]\n",
      " [11 17 16]\n",
      " [12 13 18]\n",
      " [12 18 17]\n",
      " [13 14 19]\n",
      " [13 19 18]\n",
      " [15 16 21]\n",
      " [15 21 20]\n",
      " [16 17 22]\n",
      " [16 22 21]\n",
      " [17 18 23]\n",
      " [17 23 22]\n",
      " [18 19 24]\n",
      " [18 24 23]]\n",
      "x0:\n",
      "[[ 0.    -0.     0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.75   0.   ]\n",
      " [ 0.    -1.125  0.   ]\n",
      " [ 0.    -1.5    0.   ]\n",
      " [ 0.375 -0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375 -0.75   0.   ]\n",
      " [ 0.375 -1.125  0.   ]\n",
      " [ 0.375 -1.5    0.   ]\n",
      " [ 0.75  -0.     0.   ]\n",
      " [ 0.75  -0.375  0.   ]\n",
      " [ 0.75  -0.75   0.   ]\n",
      " [ 0.75  -1.125  0.   ]\n",
      " [ 0.75  -1.5    0.   ]\n",
      " [ 1.125 -0.     0.   ]\n",
      " [ 1.125 -0.375  0.   ]\n",
      " [ 1.125 -0.75   0.   ]\n",
      " [ 1.125 -1.125  0.   ]\n",
      " [ 1.125 -1.5    0.   ]\n",
      " [ 1.5   -0.     0.   ]\n",
      " [ 1.5   -0.375  0.   ]\n",
      " [ 1.5   -0.75   0.   ]\n",
      " [ 1.5   -1.125  0.   ]\n",
      " [ 1.5   -1.5    0.   ]]\n",
      "v0:\n",
      "[[0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]]\n"
     ]
    }
   ],
   "source": [
    "# define the initial conditions: x0, v0\n",
    "# define x0\n",
    "cloth_length = 1.5\n",
    "x0, faces = dhu.standard_rectangle(cloth_length, cloth_length, Nx, Ny)\n",
    "\n",
    "# define b_u, b_v\n",
    "b_u = cloth_length / Nx\n",
    "b_v = cloth_length / Ny\n",
    "\n",
    "# define v0\n",
    "v0_0 = jnp.zeros(((N+1)**2, 3))\n",
    "#v0_1 = v0_0.at[0].set([-0.2, -0.1, 0.3])\n",
    "#v0_2 = v0_1.at[int(((N+1)**2-1)/2)].set([0, 0, 1.5])\n",
    "#v0 = v0_2.at[-1].set([0.5, 0.4, -0.1])\n",
    "v0 = v0_0\n",
    "\n",
    "# inspect faces\n",
    "print(\"faces:\")\n",
    "print(faces)\n",
    "\n",
    "# inspect x0\n",
    "print(\"x0:\")\n",
    "print(x0)\n",
    "\n",
    "# inspect v0\n",
    "print(\"v0:\")\n",
    "print(v0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7d580f41",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.2 0.  0.  ... 0.  0.  0. ]\n",
      " [0.  0.2 0.  ... 0.  0.  0. ]\n",
      " [0.  0.  0.2 ... 0.  0.  0. ]\n",
      " ...\n",
      " [0.  0.  0.  ... 0.2 0.  0. ]\n",
      " [0.  0.  0.  ... 0.  0.2 0. ]\n",
      " [0.  0.  0.  ... 0.  0.  0.2]]\n"
     ]
    }
   ],
   "source": [
    "# define the mass matrix\n",
    "# we assume all particles have the same mass m\n",
    "# 3*(N+1)^2, 3*(N+1)^2\n",
    "m_per_particle = 0.2\n",
    "M = jnp.diag(m_per_particle*jnp.ones((3*(N+1)**2)))\n",
    "\n",
    "# check M\n",
    "print(M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "951dc5ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define k_stretch, k_shear and k_bend\n",
    "k_stretch = 2.0\n",
    "k_shear = 2.0\n",
    "k_bend = 2.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ef8526c8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0\n",
      "[1, 2]\n"
     ]
    }
   ],
   "source": [
    "# for each triangle, construct u, v matrix\n",
    "# N+1, N+1\n",
    "#uu = range(0, N+1)\n",
    "#vv = range(0, N+1)\n",
    "#u, v = torch.meshgrid(uu, vv)\n",
    "#print(u)\n",
    "#print(v)\n",
    "\n",
    "# define helper functions\n",
    "def get_u(vertex_index, Ny):\n",
    "    r\"\"\"\n",
    "    Compute u from vertex index.\n",
    "    \"\"\"\n",
    "    return float(jnp.floor(vertex_index / (Ny+1)))\n",
    "\n",
    "def get_v(vertex_index, Nx):\n",
    "    r\"\"\"\n",
    "    Compute v from vertex index.\n",
    "    \"\"\"\n",
    "    return float(jnp.floor(vertex_index % (Nx+1)))\n",
    "\n",
    "def get_face_indices(vertex_index, faces):\n",
    "    r\"\"\"\n",
    "    Get a list of face indices that a vertex belongs\n",
    "    to.\n",
    "    \"\"\"\n",
    "    face_index_list = []\n",
    "    for face_index in range(faces.shape[0]):\n",
    "        face = faces[face_index]\n",
    "        if vertex_index in face:\n",
    "            face_index_list.append(face_index)\n",
    "    return face_index_list\n",
    "\n",
    "# check\n",
    "print(get_u(4, 2))\n",
    "print(get_face_indices(2, jnp.array([[0, 1, 3], [2, 9, 10], [1, 2, 5]])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "895dc839",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]\n",
      " [0. 1.]\n",
      " [1. 1.]]\n",
      "[[1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]\n",
      " [1. 1.]\n",
      " [1. 0.]]\n",
      "[0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5\n",
      " 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5]\n"
     ]
    }
   ],
   "source": [
    "# define del_u, del_v for each triangle\n",
    "# 2*N^2, 2\n",
    "# they should remain constant regardless\n",
    "# of changes to world positions\n",
    "assert 2*N**2 == faces.shape[0]\n",
    "del_u_0 = jnp.zeros((2*N**2, 2))\n",
    "del_v_0 = jnp.zeros((2*N**2, 2))\n",
    "del_us = [del_u_0]\n",
    "del_vs = [del_v_0]\n",
    "\n",
    "for face_index in range(faces.shape[0]):\n",
    "    triangle = faces[face_index]\n",
    "    # define [u, v] for point i, j, k\n",
    "    # we define the first point to be i, 2nd to be j,\n",
    "    # 3rd to be k\n",
    "    uv_i = jnp.array([get_u(triangle[0], N), get_v(triangle[0], N)])\n",
    "    uv_j = jnp.array([get_u(triangle[1], N), get_v(triangle[1], N)])\n",
    "    uv_k = jnp.array([get_u(triangle[2], N), get_v(triangle[2], N)])\n",
    "    # compute (2, ) del_u\n",
    "    del_us.append(del_us[face_index].at[face_index].set([uv_j[0]-uv_i[0], uv_k[0]-uv_i[0]]))\n",
    "    # compute (2, ) del_v\n",
    "    del_vs.append(del_vs[face_index].at[face_index].set([uv_j[1]-uv_i[1], uv_k[1]-uv_i[1]]))\n",
    "\n",
    "del_u = del_us[-1]\n",
    "del_v = del_vs[-1]\n",
    "\n",
    "# compute area of triangles\n",
    "# 2*N^2\n",
    "area_0 = jnp.zeros((2*N**2))\n",
    "areas = [area_0]\n",
    "for face_index in range(faces.shape[0]):\n",
    "    areas.append(areas[face_index].at[face_index].set(\n",
    "        abs(0.5 * (del_u[face_index, 0]*del_v[face_index, 1] - del_v[face_index, 0]*del_u[face_index, 1]))))\n",
    "\n",
    "area = areas[-1]\n",
    "\n",
    "# check del_u, del_v, area\n",
    "print(del_u)\n",
    "print(del_v)\n",
    "print(area)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "71f1f4a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.375 -0.375  0.   ]]\n",
      "[[ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]\n",
      " [ 0.375 -0.375  0.   ]\n",
      " [ 0.375  0.     0.   ]]\n",
      "[[0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]\n",
      " [0.375 0.    0.   ]]\n",
      "[[ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]\n",
      " [ 0.    -0.375  0.   ]]\n"
     ]
    }
   ],
   "source": [
    "# define del_x1\n",
    "def del_x1(x):\n",
    "    r\"\"\"\n",
    "    Compute del(x1) from x.\n",
    "    \n",
    "    :param x: ((N+1)^2, 3) array\n",
    "    \n",
    "    Return:\n",
    "    dx1: (2*N^2, 3) array\n",
    "    \"\"\"\n",
    "    # for each triangle, compute del(x1)\n",
    "    # 2*N^2, 3\n",
    "    dx1_0 = jnp.zeros((2*N**2, 3))\n",
    "    dx1s = [dx1_0]\n",
    "    \n",
    "    for face_index in range(faces.shape[0]):\n",
    "        # build del_x1: (3, )\n",
    "        dx1s.append(dx1s[face_index].at[face_index].set(\n",
    "            x[faces[face_index][1]] - x[faces[face_index][0]]))\n",
    "        \n",
    "    dx1 = dx1s[-1]\n",
    "    return dx1\n",
    "\n",
    "# define del_x2\n",
    "def del_x2(x):\n",
    "    r\"\"\"\n",
    "    Compute del(x2) from x.\n",
    "    \n",
    "    :param x: ((N+1)^2, 3) array\n",
    "    \n",
    "    Return:\n",
    "    dx2: (2*N^2, 3) array\n",
    "    \"\"\"\n",
    "    # for each triangle, compute del(x1)\n",
    "    # 2*N^2, 3\n",
    "    dx2_0 = jnp.zeros((2*N**2, 3))\n",
    "    dx2s = [dx2_0]\n",
    "    \n",
    "    for face_index in range(faces.shape[0]):\n",
    "        # build del_x1: (3, )\n",
    "        dx2s.append(dx2s[face_index].at[face_index].set(\n",
    "            x[faces[face_index][2]] - x[faces[face_index][0]]))\n",
    "        \n",
    "    dx2 = dx2s[-1]\n",
    "    return dx2\n",
    "\n",
    "# define w_u\n",
    "def dw_du(dx1, dx2):\n",
    "    r\"\"\"\n",
    "    Compute w_u from dx1 and dx2.\n",
    "    \n",
    "    :param dx1: del(x1), (2*N^2, 3) array\n",
    "    :param dx2: del(x2), (2*N^2, 3) array\n",
    "    \n",
    "    Return:\n",
    "    w_u: (2*N^2, 3) array\n",
    "    \"\"\"\n",
    "    # for each triangle, compute w_u\n",
    "    # 2*N^2, 3\n",
    "    w_u_0 = jnp.zeros((2*N**2, 3))\n",
    "    w_us = [w_u_0]\n",
    "    \n",
    "    for face_index in range(faces.shape[0]):\n",
    "        # build (2, 2) matrix del_uv = [del_u1 del_u2; del_v1 del_v2]\n",
    "        del_uv = jnp.array([del_u[face_index], del_v[face_index]])\n",
    "        # build (3, 2) matrix del_x = [dx1 dx2]\n",
    "        dx = jnp.transpose(jnp.array([dx1[face_index], dx2[face_index]]))\n",
    "        # compute (3, ) per-triangle [wu]\n",
    "        deform = jnp.matmul(dx, jnp.linalg.inv(del_uv))\n",
    "        w_us.append(w_us[face_index].at[face_index].set(deform[:, 0]))\n",
    "    \n",
    "    w_u = w_us[-1]\n",
    "    return w_u\n",
    "\n",
    "# define w_v\n",
    "def dw_dv(dx1, dx2):\n",
    "    r\"\"\"\n",
    "    Compute w_v from dx1 and dx2.\n",
    "    \n",
    "    :param dx1: del(x1), (2*N^2, 3) array\n",
    "    :param dx2: del(x2), (2*N^2, 3) array\n",
    "    \n",
    "    Return:\n",
    "    w_v: (2*N^2, 3) array\n",
    "    \"\"\"\n",
    "    # for each triangle, compute w_v\n",
    "    # 2*N^2, 3\n",
    "    w_v_0 = jnp.zeros((2*N**2, 3))\n",
    "    w_vs = [w_v_0]\n",
    "    \n",
    "    for face_index in range(faces.shape[0]):\n",
    "        # build (2, 2) matrix del_uv = [del_u1 del_u2; del_v1 del_v2]\n",
    "        del_uv = jnp.array([del_u[face_index], del_v[face_index]])\n",
    "        # build (3, 2) matrix dx = [dx1 dx2]\n",
    "        dx = jnp.transpose(jnp.array([dx1[face_index], dx2[face_index]]))\n",
    "        # compute (3, ) per-triangle [wu]\n",
    "        deform = jnp.matmul(dx, jnp.linalg.inv(del_uv))\n",
    "        w_vs.append(w_vs[face_index].at[face_index].set(deform[:, 1]))\n",
    "    \n",
    "    w_v = w_vs[-1]\n",
    "    return w_v\n",
    "\n",
    "# check del_x1, del_x2\n",
    "dx1 = del_x1(x0)\n",
    "dx2 = del_x2(x0)\n",
    "print(dx1)\n",
    "print(dx2)\n",
    "\n",
    "# check dw_du, dw_dv\n",
    "print(dw_du(dx1, dx2))\n",
    "print(dw_dv(dx1, dx2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5782eb1e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.\n",
      " 0. 0. 0. 0. 0. 0. 0. 0.]\n",
      "[[0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]]\n"
     ]
    }
   ],
   "source": [
    "# define stretch energy\n",
    "def e_stretch(x):\n",
    "    r\"\"\"\n",
    "    Compute stretch energy for each triangle\n",
    "    from world position x.\n",
    "    \n",
    "    :param x: (3*(N+1)^2, ) array\n",
    "    \n",
    "    Return:\n",
    "    energy matrix of shape (2*N^2, )\n",
    "    \"\"\"\n",
    "    # reshape x to ((N+1)^2, 3)\n",
    "    x_reshaped = jnp.reshape(x, ((N+1)**2, 3))\n",
    "    \n",
    "    # compute del(x1), del(x2)\n",
    "    # 2*N^2, 3\n",
    "    dx1 = del_x1(x_reshaped)\n",
    "    dx2 = del_x2(x_reshaped)\n",
    "    \n",
    "    # compute w_u, w_v\n",
    "    # 2*N^2, 3\n",
    "    w_u = dw_du(dx1, dx2)\n",
    "    w_v = dw_dv(dx1, dx2)\n",
    "\n",
    "    # for each triangle, compute condition matrix C_st\n",
    "    # 2*N^2, 2\n",
    "    C_st_0 = jnp.zeros((2*N**2, 2))\n",
    "    C_sts = [C_st_0]\n",
    "    for face_index in range(faces.shape[0]):\n",
    "        C_sts.append(C_sts[face_index].at[face_index].set(\n",
    "            area[face_index] * jnp.array([jnp.linalg.norm(w_u[face_index])-b_u, jnp.linalg.norm(w_v[face_index])-b_v])))\n",
    "    \n",
    "    C_st = C_sts[-1]\n",
    "\n",
    "    # for each triangle, compute energy matrix E_st\n",
    "    # 2*N^2\n",
    "    E_st_0 = jnp.zeros((2*N**2))\n",
    "    E_sts = [E_st_0]\n",
    "    for face_index in range(faces.shape[0]):\n",
    "        E_sts.append(E_sts[face_index].at[face_index].set(\n",
    "            0.5 * k_stretch * jnp.dot(C_st[face_index], C_st[face_index])))\n",
    "    E_st = E_sts[-1]\n",
    "        \n",
    "    return E_st\n",
    "\n",
    "# define f_stretch\n",
    "def f_stretch(x):\n",
    "    r\"\"\"\n",
    "    Compute stretch force on particles x\n",
    "    due to all triangles (that each particle is in).\n",
    "    \n",
    "    :param x: ((N+1)^2, 3) array\n",
    "    \n",
    "    Return:\n",
    "    Stretch force ((N+1)^2, 3)\n",
    "    \"\"\"\n",
    "    # flatten x to (3*(N+1)^2, )\n",
    "    x_flat = jnp.reshape(x, (3*(N+1)**2,))\n",
    "    \n",
    "    # compute Jacobian dE/dx\n",
    "    # 2*N^2, 3*(N+1)^2\n",
    "    dEdx = jacfwd(e_stretch)(x_flat)\n",
    "    \n",
    "    # sum up gradients over all triangles\n",
    "    # 3*(N+1)^2, \n",
    "    f_st = jnp.sum(dEdx, axis=0)\n",
    "    \n",
    "    # reshape to (N+1)^2, 3\n",
    "    f_st_reshaped = jnp.reshape(f_st, ((N+1)**2, 3))\n",
    "    return f_st_reshaped\n",
    "\n",
    "# check e_stretch\n",
    "e_st = e_stretch(x0)\n",
    "print(e_st)\n",
    "\n",
    "# check f_stretch\n",
    "f_st = f_stretch(x0)\n",
    "print(f_st)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c47b5649",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [-0.15820312  0.15820312  0.28125   ]\n",
      " [ 0.8496094  -1.3535156  -1.203125  ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 1.3535156  -0.8496094  -1.203125  ]\n",
      " [-3.2929688   3.2929688   4.25      ]\n",
      " [ 0.34570312 -0.6386719  -0.921875  ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.6386719  -0.34570312 -0.921875  ]\n",
      " [ 0.26367188 -0.26367188 -0.28125   ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]\n",
      " [ 0.          0.          0.        ]]\n"
     ]
    }
   ],
   "source": [
    "# define shear energy\n",
    "def e_shear(x):\n",
    "    r\"\"\"\n",
    "    Compute shear energy for each triangle\n",
    "    from world position x.\n",
    "    \n",
    "    :param x: (3*(N+1)^2, ) array\n",
    "    \n",
    "    Return:\n",
    "    energy matrix of shape (2*N^2, )\n",
    "    \"\"\"\n",
    "    # reshape x to ((N+1)^2, 3)\n",
    "    x_reshaped = jnp.reshape(x, ((N+1)**2, 3))\n",
    "    \n",
    "    # compute del(x1), del(x2)\n",
    "    # 2*N^2, 3\n",
    "    dx1 = del_x1(x_reshaped)\n",
    "    dx2 = del_x2(x_reshaped)\n",
    "    \n",
    "    # compute w_u, w_v\n",
    "    # 2*N^2, 3\n",
    "    w_u = dw_du(dx1, dx2)\n",
    "    w_v = dw_dv(dx1, dx2)\n",
    "    \n",
    "    # compute C_sh = area * w_u^T * w_v\n",
    "    # 2*N^2, \n",
    "    C_sh = jnp.multiply(area, jnp.sum(jnp.multiply(w_u, w_v), axis=1))\n",
    "    \n",
    "    # compute e_sh = 0.5 * k_shear * C_sh^T * C_sh\n",
    "    e_sh = 0.5 * k_shear * jnp.multiply(C_sh, C_sh)\n",
    "    return e_sh\n",
    "\n",
    "# define f_shear\n",
    "def f_shear(x):\n",
    "    r\"\"\"\n",
    "    Compute shear force on particles x\n",
    "    due to all triangles (that each particle is in).\n",
    "    \n",
    "    :param x: ((N+1)^2, 3) array\n",
    "    \n",
    "    Return:\n",
    "    Shear force ((N+1)^2, 3)\n",
    "    \"\"\"\n",
    "    # flatten x to (3*(N+1)^2, )\n",
    "    x_flat = jnp.reshape(x, (3*(N+1)**2,))\n",
    "    \n",
    "    # compute Jacobian dE/dx\n",
    "    # 2*N^2, 3*(N+1)^2\n",
    "    dEdx = jacfwd(e_shear)(x_flat)\n",
    "    \n",
    "    # sum up gradients over all triangles\n",
    "    # 3*(N+1)^2, \n",
    "    f_sh = jnp.sum(dEdx, axis=0)\n",
    "    \n",
    "    # reshape to (N+1)^2, 3\n",
    "    f_sh_reshaped = jnp.reshape(f_sh, ((N+1)**2, 3))\n",
    "    return f_sh_reshaped\n",
    "\n",
    "# check f_shear\n",
    "x1 = x0.copy()\n",
    "#x1[0] = [-0.1, -0.1, 0]\n",
    "x1[12] = [0, 0, 1.0]\n",
    "f_sh = f_shear(x1)\n",
    "print(f_sh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a2018e6f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(25, 3)\n",
      "[[ 2.9658709e+00 -2.9658709e+00  2.2244029e+00]\n",
      " [-2.1535089e-07  9.7846813e+00 -3.6692557e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [-9.7846813e+00  2.1535089e-07 -3.6692557e+00]\n",
      " [ 6.8188105e+00 -6.8188105e+00 -7.8141670e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  6.4641376e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  6.4641376e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]\n",
      " [ 0.0000000e+00  0.0000000e+00  0.0000000e+00]]\n"
     ]
    }
   ],
   "source": [
    "# define function to compute face normal\n",
    "def get_normals(x):\n",
    "    r\"\"\"\n",
    "    Compute normals at the faces.\n",
    "    \n",
    "    :param x: world positions of particles ((N+1)^2, 3)\n",
    "    \n",
    "    Return:\n",
    "    face normals of shape (2*N^2, 3)\n",
    "    \"\"\"\n",
    "    # get dx1 and dx2\n",
    "    # 2*N^2, 3\n",
    "    dx1 = del_x1(x)\n",
    "    dx2 = del_x2(x)\n",
    "    \n",
    "    # get the cross product dx1 x dx2\n",
    "    # 2*N^2, 3\n",
    "    normals = jnp.cross(dx1, dx2)\n",
    "    \n",
    "    # get the normalized normals\n",
    "    # 2*N^2, 3\n",
    "    norms = jnp.linalg.norm(normals, axis=1)\n",
    "    normals_normalized_0 = jnp.zeros(normals.shape)\n",
    "    normals_normalized_arr = [normals_normalized_0]\n",
    "    for face_index in range(faces.shape[0]):\n",
    "        normals_normalized_arr.append(normals_normalized_arr[face_index].at[face_index].set(\n",
    "            normals.at[face_index].get()/norms.at[face_index].get()\n",
    "        ))\n",
    "    normals_normalized = normals_normalized_arr[-1]\n",
    "    return normals_normalized        \n",
    "\n",
    "# define function to produce pairs of adjacent faces\n",
    "def get_pairs_of_adj_faces():\n",
    "    r\"\"\"\n",
    "    Get pairs of faces that share an edge.\n",
    "    \n",
    "    Return:\n",
    "    array of face pairs with shape (3*N^2-2*N, 2)\n",
    "    \"\"\"\n",
    "    # we take advantage of the fact that\n",
    "    # we have a square cloth with preset\n",
    "    # geometry\n",
    "    assert faces.shape[0] == 2*N**2\n",
    "    face_pairs_0 = jnp.zeros((3*N**2-2*N, 2))\n",
    "    face_pairs_arr = [face_pairs_0]\n",
    "    \n",
    "    # first we identify face pairs sharing\n",
    "    # diagonal edges    \n",
    "    for pair_index in range(N**2):\n",
    "        face_index = pair_index\n",
    "        face_pairs_arr.append(face_pairs_arr[pair_index].at[pair_index].set(\n",
    "            [2*face_index, 2*face_index+1]\n",
    "        ))\n",
    "    \n",
    "    # then we identify face pairs sharing\n",
    "    # horizontal edges\n",
    "    column_index = 0\n",
    "    for pair_index in range(N**2, N**2+N*(N-1), N-1):\n",
    "        # identify the index of the first face along a column\n",
    "        face_index = column_index * 2 * N\n",
    "        # add pairs along the column\n",
    "        for i in range(N-1):\n",
    "            pair_index_offseted = pair_index + i\n",
    "            face_pairs_arr.append(face_pairs_arr[pair_index_offseted].at[pair_index_offseted].set(\n",
    "                [face_index+2*i, face_index+2*i+3]\n",
    "            ))\n",
    "        column_index += 1\n",
    "    \n",
    "    # at last we identify face pairs sharing\n",
    "    # vertical edges\n",
    "    row_index = 0\n",
    "    for pair_index in range(N**2+N*(N-1), 3*N**2-2*N, N-1):\n",
    "        # identify the index of the first face along a row\n",
    "        face_index = row_index * 2 + 1\n",
    "        # add pairs along the row\n",
    "        for i in range(N-1):\n",
    "            pair_index_offseted = pair_index + i\n",
    "            face_pairs_arr.append(face_pairs_arr[pair_index_offseted].at[pair_index_offseted].set(\n",
    "                [face_index+(2*N)*i, face_index+(2*N)*i+(2*N-1)]\n",
    "            ))\n",
    "        row_index += 1\n",
    "    \n",
    "    face_pairs = face_pairs_arr[-1]\n",
    "    return face_pairs\n",
    "\n",
    "# find pairs of adjacent triangles\n",
    "# 2*N^2, 2\n",
    "adj_faces = get_pairs_of_adj_faces()\n",
    "\n",
    "# define bend energy\n",
    "def e_bend(x):\n",
    "    r\"\"\"\n",
    "    Compute bend energy on each pair of\n",
    "    adjacent triangles due to world position\n",
    "    x.\n",
    "    \n",
    "    :param x: (3*(N+1)^2, ) array\n",
    "    \n",
    "    Return:\n",
    "    energy matrix of shape (3*N^2-2N, )\n",
    "    \"\"\"\n",
    "    # reshape x to ((N+1)^2, 3)\n",
    "    x_reshaped = jnp.reshape(x, ((N+1)**2, 3))\n",
    "    \n",
    "    # for each pair, compute energy\n",
    "    face_normals = get_normals(x_reshaped)\n",
    "    E_be_0 = jnp.zeros((3*N**2-2*N, ))\n",
    "    E_bes = [E_be_0]\n",
    "    \n",
    "    for pair_index in range(len(adj_faces)):\n",
    "        # compute C_be for the pair\n",
    "        face_pair = adj_faces[pair_index]\n",
    "        n1 = face_normals[int(face_pair[0])]\n",
    "        n2 = face_normals[int(face_pair[1])]\n",
    "        C_be_pair = jnp.arccos(jnp.dot(n1, n2))\n",
    "        # compute E_be for the pair\n",
    "        E_bes.append(E_bes[pair_index].at[pair_index].set(0.5 * k_bend * C_be_pair**2))\n",
    "    E_be = E_bes[-1]\n",
    "    return E_be\n",
    "        \n",
    "# define f_bend\n",
    "def f_bend(x):\n",
    "    r\"\"\"\n",
    "    Compute bend force on particles x\n",
    "    due to all triangle pairs (for which the\n",
    "    shared edge is connected to the particle.)\n",
    "    \n",
    "    :param x: ((N+1)^2, 3) array\n",
    "    \n",
    "    Return:\n",
    "    Bend force ((N+1)^2, 3)\n",
    "    \"\"\"\n",
    "    # flatten x to (3*(N+1)^2, )\n",
    "    x_flat = jnp.reshape(x, (3*(N+1)**2,))\n",
    "    \n",
    "    # compute Jacobian dE/dx\n",
    "    # 3*N^2-2*N, 3*(N+1)^2\n",
    "    dEdx = jacfwd(e_bend)(x_flat)\n",
    "    \n",
    "    # sum up gradients over all triangles\n",
    "    # NOTE: we only consider non-nan gradients\n",
    "    # 3*(N+1)^2, \n",
    "    f_be = jnp.nansum(dEdx, axis=0)\n",
    "    \n",
    "    # reshape to (N+1)^2, 3\n",
    "    f_be_reshaped = jnp.reshape(f_be, ((N+1)**2, 3))\n",
    "    return f_be_reshaped\n",
    "\n",
    "# check get_normals()\n",
    "#print(get_normals(x0))\n",
    "\n",
    "# check get_pairs_of_adj_faces()\n",
    "#get_pairs_of_adj_faces()\n",
    "\n",
    "# check e_bend(x)\n",
    "x2 = x0.copy()\n",
    "#x2[0, 2] = 1.1\n",
    "#x2[1, 2] = -0.4\n",
    "#x2[5, 2] = -0.1\n",
    "#x2[6, 2] = -0.3\n",
    "#e_be = e_bend(x2)\n",
    "#print(e_be)\n",
    "x2[0] = [0, 0, 1.0]\n",
    "\n",
    "# check f_bend(x)\n",
    "f_be = f_bend(x2)\n",
    "print(f_be.shape)\n",
    "print(f_be)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "00bca43b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.\n",
      "  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.\n",
      "  0.  0. -2.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.\n",
      "  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.\n",
      "  0.  0.  0.]\n"
     ]
    }
   ],
   "source": [
    "# define external force\n",
    "def f_external(t):\n",
    "    r\"\"\"\n",
    "    Return an external force on each particle dependent\n",
    "    of time.\n",
    "    \n",
    "    :param t: time\n",
    "    \n",
    "    Return:\n",
    "    External force ((N+1)^2, 3)\n",
    "    \"\"\"\n",
    "    f_ext_0 = jnp.zeros(((N+1)**2, 3))\n",
    "    f_exts = [f_ext_0]\n",
    "    for vertex_index in range(((N+1)**2)):\n",
    "        if vertex_index == 12:\n",
    "            # we apply a constant force at the center of the cloth\n",
    "            f_exts.append(f_exts[vertex_index].at[vertex_index].set([0, 0, -2.0]))\n",
    "        else:\n",
    "            f_exts.append(f_exts[vertex_index].at[vertex_index].set([0, 0, 0]))\n",
    "    f_ext = f_exts[-1]\n",
    "    return f_ext\n",
    "\n",
    "# define f_total (total force)\n",
    "def f_total(x, t):\n",
    "    r\"\"\"\n",
    "    Compute total force on particles at position\n",
    "    x due to stretch, shear and bend.\n",
    "    \n",
    "    :param x: (3*(N+1)^2, )\n",
    "    :param t: time\n",
    "    \n",
    "    Return:\n",
    "    Force array (3*(N+1)^2, )\n",
    "    \"\"\"\n",
    "    # reshape x to ((N+1)^2, 3)\n",
    "    x_reshaped = jnp.reshape(x, ((N+1)**2, 3))\n",
    "    \n",
    "    # compute stretch\n",
    "    # (N+1)^2, 3\n",
    "    f_st = f_stretch(x_reshaped)\n",
    "    \n",
    "    # compute shear\n",
    "    # (N+1)^2, 3\n",
    "    f_sh = f_shear(x_reshaped)\n",
    "    \n",
    "    # compute bend\n",
    "    f_be = f_bend(x_reshaped)\n",
    "    \n",
    "    # compute an external force\n",
    "    # (N+1)^2, 3\n",
    "    f_ext = f_external(t)\n",
    "    \n",
    "    # add up all forces\n",
    "    f_tot = f_st + f_sh + f_be + f_ext\n",
    "    #f_tot = jnp.zeros(((N+1)**2, 3))\n",
    "    \n",
    "    # flatten total force\n",
    "    f_tot_flat = jnp.reshape(f_tot, (3*(N+1)**2, ))\n",
    "    \n",
    "    return f_tot_flat\n",
    "\n",
    "# check f_total\n",
    "f_tot = f_total(x0.reshape((3*(N+1)**2, )), 1.0)\n",
    "print(f_tot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "838ba8aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.57031244 0.         0.         ... 0.         0.         0.        ]\n",
      " [0.         0.57031244 0.         ... 0.         0.         0.        ]\n",
      " [0.         0.         0.         ... 0.         0.         0.        ]\n",
      " ...\n",
      " [0.         0.         0.         ... 0.57031244 0.         0.        ]\n",
      " [0.         0.         0.         ... 0.         0.57031244 0.        ]\n",
      " [0.         0.         0.         ... 0.         0.         0.        ]]\n"
     ]
    }
   ],
   "source": [
    "# define df/dx\n",
    "# 3*(N+1)^2, 3*(N+1)^2\n",
    "def df_dx(x, t):\n",
    "    r\"\"\"\n",
    "    Compute the total force Jacobian w.r.t. world\n",
    "    position x.\n",
    "    \n",
    "    :param x: ((N+1)^2, 3)\n",
    "    :param t: time\n",
    "    \n",
    "    Return:\n",
    "    Jacobian df/dx of shape (3*(N+1)^2, 3*(N+1)^2)\n",
    "    \"\"\"\n",
    "    # flatten x\n",
    "    x_flat = jnp.reshape(x, (3*(N+1)**2, ))\n",
    "    \n",
    "    # compute df/dx\n",
    "    dfdx = jacfwd(f_total)(x_flat, t)\n",
    "    return dfdx\n",
    "\n",
    "# check df_dx\n",
    "dfdx = df_dx(x0, 1.0)\n",
    "print(dfdx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "c124ace0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define df/dv\n",
    "# since we do not simulate damping,\n",
    "# we can assume df_dv = 0\n",
    "def df_dv():\n",
    "    return jnp.zeros((3*(N+1)**2, 3*(N+1)**2))\n",
    "\n",
    "dfdv = df_dv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "887bb154",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.000e+00  0.000e+00  0.000e+00]\n",
      " [ 0.000e+00 -3.750e-01  0.000e+00]\n",
      " [ 0.000e+00 -7.500e-01  0.000e+00]\n",
      " [ 0.000e+00 -1.125e+00  0.000e+00]\n",
      " [ 0.000e+00 -1.500e+00  0.000e+00]\n",
      " [ 3.750e-01  0.000e+00  0.000e+00]\n",
      " [ 3.750e-01 -3.750e-01  0.000e+00]\n",
      " [ 3.750e-01 -7.500e-01  0.000e+00]\n",
      " [ 3.750e-01 -1.125e+00  0.000e+00]\n",
      " [ 3.750e-01 -1.500e+00  0.000e+00]\n",
      " [ 7.500e-01  0.000e+00  0.000e+00]\n",
      " [ 7.500e-01 -3.750e-01  0.000e+00]\n",
      " [ 7.500e-01 -7.500e-01 -2.000e-04]\n",
      " [ 7.500e-01 -1.125e+00  0.000e+00]\n",
      " [ 7.500e-01 -1.500e+00  0.000e+00]\n",
      " [ 1.125e+00  0.000e+00  0.000e+00]\n",
      " [ 1.125e+00 -3.750e-01  0.000e+00]\n",
      " [ 1.125e+00 -7.500e-01  0.000e+00]\n",
      " [ 1.125e+00 -1.125e+00  0.000e+00]\n",
      " [ 1.125e+00 -1.500e+00  0.000e+00]\n",
      " [ 1.500e+00  0.000e+00  0.000e+00]\n",
      " [ 1.500e+00 -3.750e-01  0.000e+00]\n",
      " [ 1.500e+00 -7.500e-01  0.000e+00]\n",
      " [ 1.500e+00 -1.125e+00  0.000e+00]\n",
      " [ 1.500e+00 -1.500e+00  0.000e+00]]\n",
      "0.08\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD8CAYAAACxd9IeAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAQ4klEQVR4nO3dX6wc5XnH8e+vNi4pUMwf17JiFBMVgXzRGHyUgEAohZICjUguEAJFEaqIkKqkAqVSilupbaRKTXqRhosqFQJSLgiBEGiQFSVxHaKqvTDxAZMYGwqkRhhhH5vGIqVSW5OnF/MeezjZ9c7uzuzM7Pv7SEdnd874PO/u+jnvM7PvzqOIwMzm36+1PQAzmw0nu1kmnOxmmXCym2XCyW6WCSe7WSZmmuySrpf0kqRXJN3TUIwHJS1J2lvadq6kHZJeTt/PqTnmBZKelrRP0guS7mo6rqTTJT0j6fkU84tp+4WSdqXn+FFJa+qKmX7/KknPSdo+i3gpxgFJP5W0R9LutK3p13StpMclvShpv6QrGn49L06Pb/nrbUl31xozImbyBawCXgU+CKwBngc2NxDnauAyYG9p298C96Tb9wBfrjnmBuCydPss4N+BzU3GBQScmW6fBuwCLgceA25N2/8B+KOaH+vngW8A29P9RuOl33sAOH/FtqZf04eAz6Tba4C1TccsxV4FHAI+UGfM2gd6igdwBfD90v1twLaGYm1akewvARvS7Q3ASw0/1u8A180qLvAbwLPAR4CjwOpBz3kNcTYCO4FrgO3pD05j8UpxByV7Y88tcDbwH4BmFXNFnI8B/1Z3zFmW8e8HXi/dP5i2zcL6iHgz3T4ErG8qkKRNwKUUM22jcVNJvQdYAnZQVE7HIuJ42qXu5/irwBeAX6b75zUcb1kAP5C0KOnOtK3J5/ZC4Ajw9XTIcr+kMxqOWXYr8Ei6XVvM7E7QRfEnspE1wpLOBL4N3B0RbzcdNyLejYgtFDPuh4FL6vz9ZZI+DixFxGJTMU7hqoi4DLgB+Kykq8s/bOC5XU1xKPi1iLgUeIeihG4yJgDpnMdNwLdW/mzamLNM9jeAC0r3N6Zts3BY0gaA9H2p7gCSTqNI9Icj4olZxQWIiGPA0xRl9FpJq9OP6nyOrwRuknQA+CZFKX9vg/FOiIg30vcl4EmKP2xNPrcHgYMRsSvdf5wi+Wfxet4APBsRh9P92mLOMtl/DFyUzt6uoShVnppR7KeA29Pt2ymOqWsjScADwP6I+Mos4kpaJ2ltuv0+inME+ymS/ua6Y0bEtojYGBGbKF67H0bEp5qKt0zSGZLOWr5NcTy7lwaf24g4BLwu6eK06VpgX5MxS27jZAlPrTGbOMFwihMPN1KcqX4V+POGYjwCvAn8H8Vf6Dsoji13Ai8D/wycW3PMqyjKq58Ae9LXjU3GBX4HeC7F3Av8Rdr+QeAZ4BWKUvDXG3iOP8rJs/GNxku///n09cLy/5sZvKZbgN3p+f0n4JwZxDwDeAs4u7SttphKv9DM5lx2J+jMcuVkN8uEk90sE052s0w42c0yMfNkLy13dEzHdMwZxpwq2Sf8yOrMnzTHdEzHnCLZJa0C/p5ied9m4DZJm+samJnVa+JFNZKuAP4qIn4/3d8GEBF/c4p/cyLY1q1bB+6zuDj4cxbD9q/iyJEjrFu3bux/Vx7LqPgr95005jQc0zEPHDjA0aNHNehn0yT7zcD1EfGZdP/TwEci4nMr9ruTk+XIiYwZFrdYZv6r2ljpVx7LqPjj7GvWlIWFBXbv3j0wiRo/QRcR90XEQkQsbN269eQ6XenE14r9y2uFO2PQWMuPoavjNls2TbK3+ZFVMxvTNMne5kdWzWxMq0fvMlhEHJf0OeD7FBfIezAiXhjj35+4Pex4t+2SeFT8tsdnNo6Jkx0gIr4LfLemsZhZg7xc1iwTU83s41pcXDxRsg8r10e99dbVt+bKqryt6EMAmzXP7GaZcLKbZWKmZXzZpCVtlZK/S1y6W1d4ZjfLhJPdLBMzTfby2vhha8mHXvO6g+voT3H9786M0WyZZ3azTLR2gm6YUUtnq7y37ffizX6VZ3azTDjZzTLRuTK+zO/Fm9XHM7tZJpzsZpnoXBk/6QUj+noBjKr7mE3LM7tZJpzsZpnoXBk/SpVyfZ4X3vjMvU3KM7tZJkYmu6QHJS1J2lvadq6kHZJeTt/PaXaYZjatKjP7PwLXr9h2D7AzIi4Cdqb7Mzesq0wVffo0mjvPWB1GJntE/Avwnys2fwJ4KN1+CPhkvcMys7pNesy+PiLeTLcPAeuH7SjpTkm7Je0+cuTIhOHMbFpTn6CLop4cWlOWGzvW0fJ23AtGDLoARt8ugtH2+Gw+TJrshyVtAEjfl+obkpk1YdJkfwq4Pd2+HfhOPcMxs6aMXFQj6RHgo8D5kg4Cfwl8CXhM0h3Aa8AtTQ5yXHUtvOmKPo3VumtkskfEbUN+dG3NYzGzBnkFnVkmerc2flzzsJa8b1fesW7yzG6WCSe7WSbmsoyf5uowXS/7uzgm6wfP7GaZcLKbZWKmZfzi4uLAs8ldKk1Hnfnu6sKccQ8/un64YvXzzG6WCSe7WSZa689uzRn2UV5f8SZvntnNMtHa++x9mFH6dClnd56xUTyzm2XCyW6WidbK+GHlcFffux40jnnuPFPe3+X/fPDMbpYJJ7tZJlpbLuvSsJu6+m6DTc8zu1kmqjR2vEDS05L2SXpB0l1pu5s7mvVIlZn9OPAnEbEZuBz4rKTNTNDcscpy2SodX2aljs4zXX88p3pMNl+qNHZ8MyKeTbd/AewH3o+bO5r1yljH7JI2AZcCu6jY3NGNHc26oXKySzoT+DZwd0S8Xf7ZqZo7TtPYcdCnt8rbpunPXodhYxlW0lf5NFoXH1OVfa37KiW7pNMoEv3hiHgibXZzR7MeqXI2XsADwP6I+ErpR27uaNYjVRbVXAl8GvippD1p25/RUHNHL+roNr8+/VWlseO/AsMOytzc0awnvILOLBOd6wgzqjTsUuk4j51n+vT823g8s5tlwslulonOlfHj6OvVbmC88fbt8Qzb34cA7fLMbpYJJ7tZJnpXxrss7D6/Rt3kmd0sE052s0z0rozvW5ujOhbezMvjsXZ5ZjfLRO9m9nFVeS++7ZmoT++zV9G38ebCM7tZJpzsZpmYyzJ+3BK9SyX9PKiyFNhmzzO7WSac7GaZmMsyvm/vBU96wYg+HH50dVw58sxulokql5I+XdIzkp5PjR2/mLZfKGmXpFckPSppTfPDNbNJVZnZ/we4JiI+BGwBrpd0OfBl4O8i4reBnwN3NDbKlnSlS8sw89J5ps59bbgqjR0jIv4r3T0tfQVwDfB42u7GjmYdV7X906rUIGIJ2AG8ChyLiONpl4MUnV0H/Vs3djTrgErJHhHvRsQWYCPwYeCSqgGmaezYNvcqn41Rhx1+Heox1tn4iDgGPA1cAayVtPzW3UbgjXqHZmZ1qnI2fp2kten2+4DrgP0USX9z2s2NHc06rsqimg3AQ5JWUfxxeCwitkvaB3xT0l8Dz1F0ep1bXS8h+7bwpm8Ln+ZBlcaOPwEuHbD9ZxTH72bWA15BZ5aJuVwb34RB14Pr0hVZ6ug8M2yftsrpccbS1cOVLvHMbpYJJ7tZJlzGn4JLw27z6zMez+xmmXCym2XCZfwpTHoFmTbU0Xmm6u+ZlT49/33gmd0sE052s0y4jK9JlTZTZX1aeNPnxzNs/xwPATyzm2XCyW6WCZfxU8i9LOwDv0YneWY3y4Rn9in07QIMk463q7Nj357/tnlmN8uEk90sEy7jZ6iv78XXcQGMlfvMyqTvxc9j+e+Z3SwTlZM9dYV5TtL2dN+NHc16ZJyZ/S6K68Uvm/vGjnUodztxZ5NuGvQazaOqvd42An8A3J/uCzd2NOuVqjP7V4EvAL9M98/DjR3NeqVK+6ePA0sRsThJgD43dqxDldK9vE/bpf6ocQwb67DDlS4+nlM9pnlW5a23K4GbJN0InA78JnAvqbFjmt3d2NGs40bO7BGxLSI2RsQm4FbghxHxKdzY0axXpnmf/U+Bz0t6heIYfq4bO87KoF7lw/qXt6E8lmGl+7B9uqpPY53GWCvoIuJHwI/SbTd2NOsRr6Azy4TXxnfAOB8h7erHTfusyrr+eeCZ3SwTTnazTLiM74A+dT6Zx84zZV0dVx08s5tlwslulgmX8XOiS1eHGaaOK960fbWbKvG7erjimd0sE052s0y4jJ8TuSwMaVuVi4Z2qXQv88xulgknu1kmXMbPoa6WkWV9WngzL22mPLObZcLJbpYJl/FzqK9tpmC8RTV9ezzD9p/VeD2zm2XCM/uc6OKJLXuvtl8jz+xmmag0s0s6APwCeBc4HhELks4FHgU2AQeAWyLi580M08ymNc7M/rsRsSUiFtL9e4CdEXERsDPdt5bMQ+eZKt1n3HlmctOU8Z+gaOgIbuxo1nlVkz2AH0halHRn2rY+It5Mtw8B6wf9Qzd2NOuGqmfjr4qINyT9FrBD0ovlH0ZESBpYm0TEfcB9AAsLCz5N3BHjvBffVtlZx/vsfX0vvokz95Vm9oh4I31fAp6k6ARzWNKGNLANwFItIzKzRlRp2XyGpLOWbwMfA/YCT1E0dAQ3djTrvCpl/HrgyVRWrAa+ERHfk/Rj4DFJdwCvAbc0N0yrQ9uLOmy05deoiddnZLKnBo4fGrD9LeDa2kdkZo3wCjqzTHhtfEb6dhGGSTvldPVwpe3n3zO7WSac7GaZcBlvQw06M9y3RSp96jwD441l3MMVz+xmmXCym2XCZby9R1fPZFthmtfHM7tZJpzsZplwGW/vMelCljZMs0ilq4crTT7/ntnNMuFkN8uEy3ibWB96lfdtUc0o0yxq8sxulgknu1kmXMbbWMYt0btU0s+DKuv+h/HMbpYJz+w2lrYvwDCuScfbh4rEy2XNbKBKyS5praTHJb0oab+kKySdK2mHpJfT93OaHqyZTa7qzH4v8L2IuITiSrP7cWNHm0K5QWP5qyuGNZMcNtYuPoaVqjSJOBu4GngAICL+NyKO4caOZr1SZWa/EDgCfF3Sc5LuT51h3NjRrEeqJPtq4DLgaxFxKfAOK0r2KE4LDm3sGBELEbGwbt26acdrc6LtXuVN6mpJXyXZDwIHI2JXuv84RfK7saNZj4xM9og4BLwu6eK06VpgH27saNYrVRfV/DHwsKQ1wM+AP6T4Q+HGjja1rpfyfbsAxjCVkj0i9gALA37kxo5mPeEVdGaZ8Np4a12Vi2AM+vmsNNF5po3H4ZndLBNOdrNMuIy3VrRd0o5jmgtGDNPG4/fMbpYJJ7tZJlzGWyv6dsWbZcPG1IeFN57ZzTLhZDfLhMt467xxFt2s3GdW6l5408Rj8Mxulgknu1kmXMZbZzVZ0tat7oU3TZzF98xulgknu1kmXMZbZ40qX7ta3nd14Y1ndrNMONnNMuEy3ubCPC+8qetqN1XaP10saU/p621Jd7uxo1m/VLlu/EsRsSUitgBbgf8GnsSNHa1l5c4rfegwU/cYx+08M+4x+7XAqxHxGm7saNYr4yb7rcAj6bYbO5r1SOVkT91gbgK+tfJnbuxobahSFpf36VKpP2wsg8Y67HBl3Mczzsx+A/BsRBxO993Y0axHxkn22zhZwoMbO5r1SqVkl3QGcB3wRGnzl4DrJL0M/F66b9ZZw85ed6mfenksg0r3acZatbHjO8B5K7a9hRs7mvWGl8uaZcLLZW2ujbu8tO0LZizHbaJRpGd2s0w42c0y4TLe5tq8dZ4Ztk+Vs/Oe2c0y4WQ3y4RmWcJIOgK8AxydWdDC+Y7pmJnE/EBEDPwQykyTHUDS7ohYcEzHdMzZxnQZb5YJJ7tZJtpI9vsc0zEdc/YxZ37MbmbtcBlvlgknu1kmnOxmmXCym2XCyW6Wif8HgV6wyb0B9KIAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# define v and x update\n",
    "def compute_dv(v_curr, x_curr, dt, t, viz=False):\n",
    "    r\"\"\"\n",
    "    Compute 1) the change in velocity dv given the\n",
    "    current velocity, position and time step; 2) error\n",
    "    between dv computed by np.linalg.inv() and dv\n",
    "    computed via MPCG.\n",
    "    \n",
    "    :param v_curr: current velocity, ((N+1)^2, 3)\n",
    "    :param x_curr: current position, ((N+1)^2, 3)\n",
    "    :param dt: time step size, scalar\n",
    "    :param t: time\n",
    "    :param viz: visualization of solver(s), on/off\n",
    "    \n",
    "    Return:\n",
    "    dv, change in velocity, ((N+1)^2, 3)\n",
    "    \"\"\"\n",
    "    # flatten v_curr\n",
    "    v_curr_flat = jnp.reshape(v_curr, (3*(N+1)**2, ))\n",
    "    \n",
    "    # flatten x_curr\n",
    "    x_curr_flat = jnp.reshape(x_curr, (3*(N+1)**2, ))\n",
    "    \n",
    "    # we compute df_dx and save it\n",
    "    dfdx = df_dx(x_curr, t)\n",
    "    \n",
    "    # compute dv with MPCG solver:\n",
    "    # we formulate the problem as A*dv = b \n",
    "    # compute A = M - h*df/dv - h^2*df/dx\n",
    "    A = jnp.eye(3*(N+1)**2) - dt * dfdv - dt**2 * dfdx\n",
    "    # visualize the sparcity of A\n",
    "    if viz:\n",
    "        plt.spy(A)\n",
    "        plt.savefig('viz_A_sparcity.png')\n",
    "    # compute b = h*(f+h*df/dx*v)\n",
    "    b = dt * (f_total(x_curr_flat, t) + dt * jnp.matmul(dfdx, v_curr_flat))\n",
    "    # create the constraint list\n",
    "    S = []\n",
    "    for particle_index in range((N+1)**2):\n",
    "        S.append(())\n",
    "    # fix the four edges\n",
    "    constraint_dirs = (\n",
    "        np.array([1.0, 0, 0]),\n",
    "        np.array([0.0, 1, 0]),\n",
    "        np.array([0.0, 0, 1]),\n",
    "    )\n",
    "    S[1] = constraint_dirs\n",
    "    S[4] = constraint_dirs\n",
    "    S[20] = constraint_dirs\n",
    "    S[24] = constraint_dirs\n",
    "    z = np.zeros(((N+1)**2, 3))\n",
    "    # solve with MPCG\n",
    "    mpcg_solver = MPCGSolver(A.copy(), b.copy(), S, z)\n",
    "    dv_mpcg_2d = mpcg_solver.solve()\n",
    "    # remove the extra dimension in dv_mpcg\n",
    "    # 3*(N+1)^2,\n",
    "    dv_mpcg_1d = jnp.array(dv_mpcg_2d[:, 0])\n",
    "    # reshape dv_mpcg_1d to flattend vector\n",
    "    # (N+1)^2, 3\n",
    "    dv_mpcg = jnp.reshape(dv_mpcg_1d, ((N+1)**2, 3))\n",
    "    \n",
    "    # compute dv with jnp.linalg.inv():\n",
    "    # compute the W matrix\n",
    "    S_blocks = mpcg_solver.compute_S(S)\n",
    "    W = np.zeros((3*(N+1)**2, 3*(N+1)**2))\n",
    "    for particle_index in range((N+1)**2):\n",
    "        row_start = 3 * particle_index\n",
    "        col_start = row_start\n",
    "        W[row_start:row_start+3, col_start:col_start+3] = S_blocks[particle_index]\n",
    "    W = 1.0/m_per_particle * W\n",
    "    # compute the modified A matrix\n",
    "    # A_modified = I - hWdf/dv - h^2Wdf/dx\n",
    "    A_modified = np.eye(3*(N+1)**2) - dt * np.matmul(W, dfdv.copy()) - dt**2 * np.matmul(W, dfdx.copy())\n",
    "    # compute the modified b vector\n",
    "    # b_modified = h*W*(f+h*df/dx*v) + z\n",
    "    b_modified = dt * np.matmul(W, (f_total(x_curr_flat, t)\n",
    "                                    + dt * np.matmul(dfdx, v_curr_flat))) + np.reshape(z, (3*(N+1)**2, ))\n",
    "    # get dv_flat_lu\n",
    "    # 3*(N+1)^2,\n",
    "    dv_lu_1d = np.matmul(np.linalg.inv(A_modified), b_modified)\n",
    "    # reshape dv to flattend vector\n",
    "    # (N+1)^2, 3\n",
    "    dv_lu = jnp.reshape(dv_lu_1d, ((N+1)**2, 3))\n",
    "    \n",
    "    # compute error (inf norm of ||dv_lu - dv_mpcg||)\n",
    "    error = np.linalg.norm((dv_lu_1d - dv_mpcg_1d), np.inf)\n",
    "    \n",
    "    return dv_mpcg, error\n",
    "\n",
    "def compute_v_next(v_curr, dv):\n",
    "    r\"\"\"\n",
    "    Compute the next velocity given the current\n",
    "    velocity and the velocity increment.\n",
    "    \n",
    "    :param v_curr: current velocity, ((N+1)^2, 3)\n",
    "    :param dv: velocity change, ((N+1)^2, 3)\n",
    "    \n",
    "    Return:\n",
    "    next velocity, ((N+1)^2, 3)\n",
    "    \"\"\"\n",
    "    return (v_curr + dv)\n",
    "\n",
    "def compute_x_next(x_curr, v_next, dt):\n",
    "    r\"\"\"\n",
    "    Compute the new position given the current\n",
    "    position, the next velocity and dt.\n",
    "    \n",
    "    :param x_curr: current position, ((N+1)^2, 3)\n",
    "    :param v_next: next velocity, ((N+1)^2, 3)\n",
    "    \n",
    "    Return:\n",
    "    next position, ((N+1)^2, 3)\n",
    "    \"\"\"\n",
    "    return (x_curr + dt*v_next)\n",
    "\n",
    "# check time stepping\n",
    "dt = 1e-2\n",
    "t = dt\n",
    "dv, err = compute_dv(v0, x0, dt, t, viz=True)\n",
    "v1 = compute_v_next(v0, dv)\n",
    "x1 = compute_x_next(x0, v1, dt)\n",
    "print(x1)\n",
    "print(err)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "215294da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "error: 0.1599999964237213\n",
      "error: 0.5437749624252319\n",
      "error: 0.7500970363616943\n",
      "error: 1.6328060626983643\n",
      "error: 2.5724539756774902\n",
      "error: 4.005772590637207\n",
      "error: 5.817660808563232\n",
      "error: 7.177479267120361\n",
      "error: 6.830315589904785\n",
      "error: 4.941368103027344\n",
      "error: 4.226620674133301\n",
      "error: 3.3877410888671875\n",
      "error: 2.4397130012512207\n",
      "error: 1.958341360092163\n",
      "error: 2.249234199523926\n",
      "error: 3.0066699981689453\n",
      "error: 3.9674758911132812\n",
      "error: 3.7065277099609375\n",
      "error: 3.2077460289001465\n",
      "error: 2.389683246612549\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c7b38065e3ae4cdcaecd9381eb1c4c61",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(aspect=1.5, position=(5.0, 3.0, 5.0), projectionMatrix=(1.0, 0.0, 0.0, 0.0, …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "76b39155930549e8ab08c0d504c33d0b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "AnimationAction(clip=AnimationClip(tracks=(NumberKeyframeTrack(name='.morphTargetInfluences', times=array([0. …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# define total number of steps, time step size\n",
    "Nt = 20\n",
    "dt = 2e-2\n",
    "times = []\n",
    "\n",
    "# define positions and velocities\n",
    "vt = [v0]\n",
    "xt = [x0]\n",
    "\n",
    "for k in range(0, Nt, 1):\n",
    "    x = x0.copy()\n",
    "    t = k * dt\n",
    "    # update velocity and world position,\n",
    "    # given the relation\n",
    "    # v_(k+1) = v_k + dv\n",
    "    # x_(k+1) = x_k = dt*v_(k+1)\n",
    "    dv, err = compute_dv(vt[k], xt[k], dt, t)\n",
    "    vt.append(compute_v_next(vt[k], dv))\n",
    "    x_next = compute_x_next(xt[k], vt[k+1], dt)\n",
    "    xt.append(x_next)\n",
    "    \n",
    "    # apply positional constraints (np.linalg.inv())\n",
    "    #x_next_constrained_1 = x_next.at[0].set(x0[0])\n",
    "    #x_next_constrained_4 = x_next_constrained_1.at[4].set(x0[4])\n",
    "    #x_next_constrained_20 = x_next_constrained_4.at[20].set(x0[20])\n",
    "    #x_next_constrained_24 = x_next_constrained_20.at[24].set(x0[24])\n",
    "    #xt.append(x_next_constrained_24)\n",
    "    \n",
    "    print(f\"error: {err}\")\n",
    "    \n",
    "    times.append(t)\n",
    "\n",
    "dhu.mesh_animation(times, xt, faces)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "659fcf36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "30726fa8d8964a4c92b4390c0d403e39",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(aspect=1.5, position=(5.0, 3.0, 5.0), projectionMatrix=(1.0, 0.0, 0.0, 0.0, …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# display mesh at a specific time step\n",
    "step_count = 20\n",
    "dhu.mesh_display(xt[step_count], faces)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a346d3a7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:math607e] *",
   "language": "python",
   "name": "conda-env-math607e-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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}