examples/tutorial/physics/simplicits_low_level_api.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Simplicits Full-Feature API\n",
    "Users with a physics simulation background might find it useful to use the [Simplicits](https://research.nvidia.com/labs/toronto-ai/simplicits/) API in a more customizable way. \n",
    "In this notebook, we expose the inner details of our mesh-free, geometry-agnostic elastic simulator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Notebook requires k3d\n",
    "!pip install k3d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an Object\n",
    "Use an SDF to sample points inside the object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os, sys\n",
    "import logging\n",
    "import k3d\n",
    "from functools import partial\n",
    "\n",
    "import torch \n",
    "import numpy as np \n",
    "import matplotlib.colors as mcolors\n",
    "\n",
    "from create_example_object import *\n",
    "\n",
    "# Import kaolin physics\n",
    "import kaolin.physics as physics\n",
    "print(dir(physics.simplicits.network))\n",
    "# Imports for displaying the object\n",
    "from scipy.spatial import Delaunay\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display\n",
    "\n",
    "device = 'cuda'\n",
    "dtype = torch.float32\n",
    "\n",
    "logging.basicConfig(level=logging.INFO, stream=sys.stdout)\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "so_pts, _, so_yms, so_prs, so_rhos, so_appx_vol = example_unit_cube_object()\n",
    "so_pts[:,1] += 1\n",
    "print(so_appx_vol)\n",
    "\n",
    "plot = k3d.plot()\n",
    "plot += k3d.points(so_pts.cpu().detach().numpy(), point_size=0.01)\n",
    "plot.display()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Object\n",
    "Train the object using the following training parameters. Notice that the sum of `le` or elastic loss term and the `lo` or orthogonality loss terms stabilize and converge to a fairly small value. The magnitudes of the elastic and orthogonality losses start off very different and over the training become similar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "NUM_HANDLES = 5\n",
    "NUM_STEPS = 10000\n",
    "LR_START = 1e-3\n",
    "NUM_SAMPLES = 1000\n",
    "\n",
    "ENERGY_INTERP_LINSPACE = np.linspace(0, 1, NUM_STEPS, endpoint=False)\n",
    "\n",
    "so_pts, _, so_yms, so_prs, so_rhos, so_appx_vol = example_unit_cube_object()\n",
    "so_model = physics.simplicits.network.SimplicitsMLP(spatial_dimensions=3, layer_width=64, num_handles=NUM_HANDLES, num_layers=8)\n",
    "so_optimizer = torch.optim.Adam(so_model.parameters(), LR_START)\n",
    "\n",
    "# Don't normalize object to unit cube unless is very small or very lage\n",
    "# so_bb_pts = (torch.min(so_pts, axis=0).values, torch.max(so_pts, axis=0).values)\n",
    "# so_normalized_pts = (so_pts - so_bb_pts[0])/(so_bb_pts[1] - so_bb_pts[0])\n",
    "\n",
    "so_pts = torch.as_tensor(so_pts, device=device, dtype=dtype)\n",
    "so_yms = torch.as_tensor(so_yms, device=device, dtype=dtype).unsqueeze(-1)\n",
    "so_prs = torch.as_tensor(so_prs, device=device, dtype=dtype).unsqueeze(-1)\n",
    "so_rhos = torch.as_tensor(so_rhos, device=device, dtype=dtype).unsqueeze(-1)\n",
    "so_normalized_pts = so_pts.clone()\n",
    "so_model.to(device)\n",
    "\n",
    "partial_compute_losses = partial(physics.simplicits.losses.compute_losses, \n",
    "                             batch_size=10, \n",
    "                             num_handles=NUM_HANDLES, \n",
    "                             appx_vol=1, \n",
    "                             num_samples=NUM_SAMPLES, \n",
    "                             le_coeff=1e-1, \n",
    "                             lo_coeff=1e6)\n",
    "\n",
    "\n",
    "\n",
    "so_model.train()\n",
    "for i in range(NUM_STEPS):\n",
    "    #Set grads to zero\n",
    "    so_optimizer.zero_grad()\n",
    "    #train a step\n",
    "    le, lo = partial_compute_losses(so_model, so_normalized_pts, so_yms, so_prs, so_rhos, float(i/NUM_STEPS))\n",
    "    loss = le + lo\n",
    "    # Backprop over the losses\n",
    "    loss.backward()\n",
    "    # Take optimizer step\n",
    "    so_optimizer.step()\n",
    "    \n",
    "    if i%100 == 0:\n",
    "        print(f'Training step: {i}, le: {le.item()}, lo: {lo.item()}')\n",
    "\n",
    "so_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize Learned Weights\n",
    "Visualizes the various learned skinning eigen-modes/subspaces/bases of the trained object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def scalar_to_rgb(s_normalized, mincolorstr, maxcolorstr, minval=None, maxval=None):\n",
    "    cmap = mcolors.LinearSegmentedColormap.from_list(\"\", [mincolorstr, maxcolorstr])\n",
    "\n",
    "    if(np.sum(s_normalized)==0):\n",
    "        return cmap(s_normalized)[:, 0:3]*255\n",
    "    else:\n",
    "        if minval == None:\n",
    "            minval = np.min(s_normalized)\n",
    "        if maxval == None:\n",
    "            maxval = np.max(s_normalized)\n",
    "        s_normalized = (s_normalized - np.min(s_normalized)) / (maxval - minval)\n",
    "        colors_rgb = cmap(s_normalized)[:, 0:3]*255\n",
    "        return colors_rgb\n",
    "\n",
    "def rgb2hex(rgb):\n",
    "    str_res = \"0x{0:02x}{1:02x}{2:02x}\".format(int(rgb[0]), int(rgb[1]), int(rgb[2]))\n",
    "    return int(str_res, 16)\n",
    "\n",
    "modes = so_model(so_pts)\n",
    "\n",
    "def visualize_mode(viz_mode, modes):\n",
    "    rgb = scalar_to_rgb(modes[:,viz_mode].detach().cpu().numpy(), \"blue\", \"red\")\n",
    "    colors = [rgb2hex(rgb[xx, :]) for xx in range(rgb.shape[0])]\n",
    "    return colors\n",
    "\n",
    "# Viewing 3 of the weight functions\n",
    "# Plot more to view more\n",
    "plot = k3d.plot()\n",
    "plot += k3d.points(so_normalized_pts.detach().cpu().numpy(), colors=visualize_mode(0,modes),  point_size=0.01)\n",
    "plot += k3d.points(so_normalized_pts.detach().cpu().numpy(), colors=visualize_mode(4,modes),  point_size=0.01)\n",
    "plot += k3d.points(so_normalized_pts.detach().cpu().numpy(), colors=visualize_mode(2,modes),  point_size=0.01)\n",
    "plot.display()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulate Object\n",
    "Set up the simulator with the following parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sim parameters\n",
    "NUM_STEPS = 100\n",
    "DT = 0.05\n",
    "FLOOR_PLANE = -1\n",
    "PENALTY_WEIGHT = 10000\n",
    "NUM_SAMPLES = 1000\n",
    "device = 'cuda'\n",
    "dtype = torch.float32\n",
    "MAX_NEWTON_STEPS=5\n",
    "MAX_LS_STEPS = 30\n",
    "\n",
    "# Select the cubature points and corresponding material parameters\n",
    "sample_indices = torch.randint(low=0, high=so_normalized_pts.shape[0], size=(NUM_SAMPLES,), device=so_normalized_pts.device)\n",
    "sim_pts = so_pts[sample_indices]\n",
    "sim_normalized_pts = so_normalized_pts[sample_indices]\n",
    "sim_yms = 0.1*so_yms[sample_indices]\n",
    "sim_prs = so_prs[sample_indices]\n",
    "sim_rhos = 5*so_rhos[sample_indices]\n",
    "sim_weights = torch.cat((so_model(sim_normalized_pts), torch.ones((sim_normalized_pts.shape[0], 1), device=device)), dim=1)\n",
    "model_plus_rigid = lambda pts: torch.cat((so_model(pts), torch.ones((pts.shape[0], 1), device=device)), dim=1)\n",
    "\n",
    "# Initialize simulation DOFs\n",
    "z = torch.zeros(sim_weights.shape[1]*12 , dtype=dtype, device = device).unsqueeze(-1)\n",
    "z_prev = z.clone().detach()\n",
    "z_dot = torch.zeros_like(z, device=device)\n",
    "x0_flat = sim_pts.flatten().unsqueeze(-1)\n",
    "\n",
    "\n",
    "M, invM = physics.simplicits.precomputed.lumped_mass_matrix(sim_rhos, so_appx_vol, dim = 3)\n",
    "dFdz = physics.simplicits.precomputed.jacobian_dF_dz(model_plus_rigid, sim_normalized_pts, z).detach()\n",
    "dxdz = torch.autograd.functional.jacobian(lambda x: physics.simplicits.utils.weight_function_lbs(sim_pts, tfms = x.reshape(-1,3,4).unsqueeze(0), fcn = model_plus_rigid).flatten(), z.flatten())\n",
    "bigI = torch.tile(torch.eye(3, device=device).flatten().unsqueeze(dim=1), (NUM_SAMPLES,1)).detach()\n",
    "B = physics.simplicits.precomputed.lbs_matrix(sim_pts, sim_weights).detach()\n",
    "\n",
    "# 3*num samples gravities per sample point\n",
    "grav = torch.tensor([0, 9.8, 0], device=device)\n",
    "\n",
    "BMB = B.T @ M @ B\n",
    "BinvMB = B.T @ invM @ B\n",
    "\n",
    "print(\" Density: \",str(sim_rhos[0].item())+\"kg/m^3\\n\", \n",
    "      \"Youngs Mod: \", str(sim_yms[0].item())+\"Pa\\n\", \n",
    "      \"Poiss Ratio: \", str(sim_prs[0].item())+\"\\n\", \n",
    "      \"Appx Vol: \", str(so_appx_vol)+\"m^3\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#########################SETUP MATERIAL AND SCENE FORCES######################################################\n",
    "mus, lams = physics.materials.utils.to_lame(sim_yms, sim_prs)\n",
    "material_object = physics.materials.NeohookeanMaterial(sim_yms, sim_prs)\n",
    "gravity_object = physics.utils.Gravity(rhos=sim_rhos, acceleration=grav)\n",
    "floor_object = physics.utils.Floor(floor_height=FLOOR_PLANE, floor_axis=1)\n",
    "bdry_cond = physics.utils.Boundary()\n",
    "bdry_indx = torch.nonzero(sim_pts[:,1]>1.45, as_tuple=False).squeeze()\n",
    "bdry_pos = sim_pts[bdry_indx,:]\n",
    "bdry_cond.set_pinned_verts(bdry_indx, bdry_pos)\n",
    "integration_sampling = torch.as_tensor(so_appx_vol/NUM_SAMPLES, device=device, dtype=sim_pts.dtype)\n",
    "\n",
    "#######################Physics Energy, Forces, Hessians########################################################\n",
    "partial_bdry_e = physics.simplicits.simplicits_scene_forces.generate_fcn_simplicits_scene_energy(bdry_cond, B, coeff=PENALTY_WEIGHT, integration_sampling=None)\n",
    "partial_bdry_g = physics.simplicits.simplicits_scene_forces.generate_fcn_simplicits_scene_gradient(bdry_cond, B, coeff=PENALTY_WEIGHT, integration_sampling=None)\n",
    "partial_bdry_h = physics.simplicits.simplicits_scene_forces.generate_fcn_simplicits_scene_hessian(bdry_cond, B, coeff=PENALTY_WEIGHT, integration_sampling=None)\n",
    "\n",
    "partial_grav_e = physics.simplicits.simplicits_scene_forces.generate_fcn_simplicits_scene_energy(gravity_object, B, coeff=1, integration_sampling=integration_sampling) \n",
    "partial_grav_g = physics.simplicits.simplicits_scene_forces.generate_fcn_simplicits_scene_gradient(gravity_object, B, coeff=1, integration_sampling=integration_sampling) \n",
    "partial_grav_h = physics.simplicits.simplicits_scene_forces.generate_fcn_simplicits_scene_hessian(gravity_object, B, coeff=1, integration_sampling=integration_sampling) \n",
    "\n",
    "partial_material_e = physics.simplicits.simplicits_scene_forces.generate_fcn_simplicits_material_energy(material_object, dFdz, coeff=1, integration_sampling=integration_sampling)\n",
    "partial_material_g = physics.simplicits.simplicits_scene_forces.generate_fcn_simplicits_material_gradient(material_object, dFdz, coeff=1, integration_sampling=integration_sampling)\n",
    "partial_material_h = physics.simplicits.simplicits_scene_forces.generate_fcn_simplicits_material_hessian(material_object, dFdz, coeff=1, integration_sampling=integration_sampling)\n",
    "\n",
    "partial_floor_e = physics.simplicits.simplicits_scene_forces.generate_fcn_simplicits_scene_energy(floor_object, B, coeff=PENALTY_WEIGHT, integration_sampling=None)\n",
    "partial_floor_g = physics.simplicits.simplicits_scene_forces.generate_fcn_simplicits_scene_gradient(floor_object, B, coeff=PENALTY_WEIGHT, integration_sampling=None)\n",
    "partial_floor_h = physics.simplicits.simplicits_scene_forces.generate_fcn_simplicits_scene_hessian(floor_object, B, coeff=PENALTY_WEIGHT, integration_sampling=None)\n",
    "###############################################################################\n",
    "\n",
    "####################Backwards Euler Functions###########################################################\n",
    "def potential_sum(output, z, z_dot, B, dFdz, x0_flat, bigI, defo_grad_fcns = [], pt_wise_fcns = []):\n",
    "    # updates the quantity calculated in the output value\n",
    "    F_ele = torch.matmul(dFdz, z) + bigI\n",
    "    x_flat = B @ z + x0_flat\n",
    "    x = x_flat.reshape(-1,3)\n",
    "    for e in defo_grad_fcns:\n",
    "        output += e(F_ele)\n",
    "    for e in pt_wise_fcns:\n",
    "        output += e(x)\n",
    "        \n",
    "def newton_E(z, z_prev, z_dot, B, BMB, dt, x0_flat, dFdz, bigI, defo_grad_energies = [], pt_wise_energies = []):\n",
    "    pe_sum = torch.tensor([0], device=device, dtype=dtype)\n",
    "    potential_sum(pe_sum, z, z_dot, B, dFdz, x0_flat, bigI, defo_grad_energies, pt_wise_energies)\n",
    "    return 0.5 * z.T @ BMB @ z - z.T @ BMB @ z_prev - dt * z.T @ BMB @ z_dot + dt * dt * pe_sum\n",
    "\n",
    "def newton_G(z, z_prev, z_dot, B, BMB, dt, x0_flat, dFdz, bigI, defo_grad_gradients = [], pt_wise_gradients = []):\n",
    "    pe_grad_sum = torch.zeros_like(z)\n",
    "    potential_sum(pe_grad_sum, z, z_dot, B, dFdz, x0_flat, bigI, defo_grad_gradients, pt_wise_gradients)\n",
    "    return BMB @ z - BMB @ z_prev - dt * BMB @ z_dot + dt * dt * pe_grad_sum\n",
    "\n",
    "def newton_H(z, z_prev, z_dot, B, BMB, dt, x0_flat, dFdz, bigI, defo_grad_hessians = [], pt_wise_hessians = []):\n",
    "    pe_hess_sum = torch.zeros(z.shape[0], z.shape[0], device=device, dtype=dtype)\n",
    "    potential_sum(pe_hess_sum, z, z_dot, B, dFdz, x0_flat, bigI, defo_grad_hessians, pt_wise_hessians)\n",
    "    return BMB  + dt * dt * pe_hess_sum\n",
    "    \n",
    "##########################Backwards Euler Partials#####################################################\n",
    "partial_newton_E = partial(newton_E, \n",
    "                           B=B.detach(), \n",
    "                           BMB = BMB.detach(), \n",
    "                           dt=DT, \n",
    "                           x0_flat=x0_flat.detach(), \n",
    "                           dFdz=dFdz.detach(), \n",
    "                           bigI=bigI.detach(),\n",
    "                           defo_grad_energies=[partial_material_e],\n",
    "                           pt_wise_energies=[partial_grav_e, partial_floor_e, partial_bdry_e])\n",
    "partial_newton_G = partial(newton_G, \n",
    "                           B=B.detach(), \n",
    "                           BMB = BMB.detach(), \n",
    "                           dt=DT, \n",
    "                           x0_flat=x0_flat.detach(), \n",
    "                           dFdz=dFdz.detach(), \n",
    "                           bigI=bigI.detach(),\n",
    "                           defo_grad_gradients=[partial_material_g],\n",
    "                           pt_wise_gradients=[partial_grav_g, partial_floor_g, partial_bdry_g])\n",
    "partial_newton_H = partial(newton_H, \n",
    "                           B=B.detach(), \n",
    "                           BMB = BMB.detach(), \n",
    "                           dt=DT, \n",
    "                           x0_flat=x0_flat.detach(), \n",
    "                           dFdz=dFdz.detach(), \n",
    "                           bigI=bigI.detach(),\n",
    "                           defo_grad_hessians=[partial_material_h],\n",
    "                           pt_wise_hessians=[partial_grav_h, partial_floor_h, partial_bdry_h])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation Loop\n",
    "Here's the simulation loop exposed. At each backwards euler sim step, we run newton's method to find the optimal DOFs `z`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "z = torch.zeros(sim_weights.shape[1]*12 , dtype=dtype, device = device).unsqueeze(-1)\n",
    "z_prev = z.clone().detach()\n",
    "z_dot = torch.zeros_like(z, device=device).detach()\n",
    "x0_flat = sim_pts.flatten().unsqueeze(-1)\n",
    "states = [z.clone().detach()]\n",
    "for time_step in range(int(NUM_STEPS)):\n",
    "    print(\"Timestep: \", time_step)\n",
    "    z_prev = z.clone().detach()\n",
    "    more_partial_newton_E = partial(partial_newton_E, z_prev=z_prev, z_dot=z_dot)\n",
    "    more_partial_newton_G = partial(partial_newton_G, z_prev=z_prev, z_dot=z_dot)\n",
    "    more_partial_newton_H = partial(partial_newton_H, z_prev=z_prev, z_dot=z_dot)\n",
    "    z = physics.utils.optimization.newtons_method(z, more_partial_newton_E, more_partial_newton_G, more_partial_newton_H)\n",
    "    F_ele = torch.matmul(dFdz, z) + bigI\n",
    "    x_pts = (B @ z + x0_flat).reshape(-1,3)\n",
    "    print(f'\\tFloor E:{partial_floor_e(x_pts).item()}, Grav E:{partial_grav_e(x_pts).item()}, Bdry E:{partial_bdry_e(x_pts).item()}, Elastic E:{partial_material_e(F_ele).item()}')\n",
    "    with torch.no_grad():\n",
    "        z_dot = (z - z_prev)/DT\n",
    "    states.append(z.clone().detach())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Displaying The Simulation \n",
    "Next we use K3D to display the simulation.\n",
    "Scrub through the `TIMESTEP` slider below to see each step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_ground_plane(plane):\n",
    "    # Define the size of the ground plane\n",
    "    x_size = 1\n",
    "    y_size = 1\n",
    "    num_points_x = 30\n",
    "    num_points_y = 30\n",
    "    \n",
    "    # Generate grid points\n",
    "    x = np.linspace(-x_size/2, x_size/2, num_points_x)\n",
    "    y = np.linspace(-y_size/2, y_size/2, num_points_y)\n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    \n",
    "    # Generate vertices\n",
    "    vertices = np.column_stack([X.flatten(), plane * np.ones_like(X.flatten()), Y.flatten()])\n",
    "    vertices += +1e-3*np.random.rand(vertices.shape[0], vertices.shape[1])\n",
    "    # Perform Delaunay triangulation\n",
    "    tri = Delaunay(vertices)\n",
    "    return torch.tensor(vertices), torch.tensor(tri.simplices)\n",
    "\n",
    "v, f = generate_ground_plane(FLOOR_PLANE)\n",
    "\n",
    "# Function to create points\n",
    "def create_points(t):\n",
    "    z = states[int(t)]\n",
    "    x = B @ z + x0_flat\n",
    "    # print(torch.norm(x - x0_flat))\n",
    "    scene_verts = x.reshape(-1,3).unsqueeze(0).cpu().detach()\n",
    "    return scene_verts.cpu().detach().numpy()\n",
    "\n",
    "# Create a plot\n",
    "plot = k3d.plot(camera_auto_fit=False)\n",
    "floor_plot = k3d.points(v.cpu().detach().numpy(), point_size=0.02, color=0x00ff00, shader='3d')\n",
    "plot += floor_plot\n",
    "# Initial set of points\n",
    "initial_points = create_points(0)\n",
    "points_plot = k3d.points(initial_points, point_size=0.1, color=0x00ff00, shader='3d')\n",
    "plot += points_plot\n",
    "\n",
    "# Rotate the camera so that y is the vertical axis\n",
    "plot.camera = [3, 1, 0, 0, 0, 0, 0, 1, 0]  # (eye_x, eye_y, eye_z, target_x, target_y, target_z, up_x, up_y, up_z)\n",
    "\n",
    "\n",
    "# Display the plot\n",
    "plot.display()\n",
    "\n",
    "# Define the function to update the points\n",
    "def update_points(TIMESTEP):\n",
    "    new_points = create_points(TIMESTEP)\n",
    "    points_plot.positions = new_points.astype(np.float32)\n",
    "\n",
    "# Create a slider\n",
    "slider = widgets.FloatSlider(min=0, max=len(states), step=1, value=0)\n",
    "\n",
    "# Link the slider to the update function\n",
    "widgets.interactive(update_points, TIMESTEP=slider)\n",
    "\n",
    "# Display the slider\n",
    "display(slider)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}