fix bug in test

Author: Huangzizhou

test/test_differentiable.ipynb

@@ -6,9 +6,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import sys\n",
+    "sys.path.append(\"/Users/zizhouhuang/Desktop/polyfem-python/build/\")\n",
     "import polyfempy as pf\n",
     "import json\n",
-    "import numpy as np\n",
     "import torch\n",
     "\n",
     "torch.set_default_dtype(torch.float64)"
@@ -20,33 +21,36 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Differentiable simulator that computes shape derivatives\n",
+    "# Differentiable simulator that computes initial derivatives\n",
     "class Simulate(torch.autograd.Function):\n",
     "\n",
     "    @staticmethod\n",
-    "    def forward(ctx, solver, vertices):\n",
+    "    def forward(ctx, solver, body_ids, initial_velocities):\n",
     "        # Update solver setup\n",
-    "        solver.mesh().set_vertices(vertices)\n",
+    "        for bid, vel in zip(body_ids, initial_velocities):\n",
+    "            print(bid, vel)\n",
+    "            solver.set_initial_velocity(bid, vel)\n",
+    "        sys.stdout.flush()\n",
     "        # Enable caching intermediate variables in the simulation, which will be used for solve_adjoint\n",
     "        solver.set_cache_level(pf.CacheLevel.Derivatives)\n",
     "        # Run simulation\n",
     "        solver.solve()\n",
-    "        # Collect transient simulation solutions\n",
-    "        cache = solver.get_solution_cache()\n",
-    "        sol = torch.zeros((solver.ndof(), cache.size()))\n",
-    "        for t in range(cache.size()):\n",
-    "            sol[:, t] = torch.tensor(cache.solution(t))\n",
     "        # Cache solver for backward gradient propagation\n",
     "        ctx.solver = solver\n",
-    "        return sol\n",
+    "        ctx.bids = body_ids\n",
+    "        return torch.tensor(solver.get_solutions())\n",
     "\n",
     "    @staticmethod\n",
     "    @torch.autograd.function.once_differentiable\n",
     "    def backward(ctx, grad_output):\n",
     "        # solve_adjoint only needs to be called once per solver, independent of number of types of optimization variables\n",
-    "        ctx.solver.solve_adjoint(grad_output)\n",
-    "        # Compute shape derivatives\n",
-    "        return None, torch.tensor(pf.shape_derivative(ctx.solver))"
+    "        ctx.solver.solve_adjoint(grad_output.detach().numpy())\n",
+    "        # Compute initial derivatives\n",
+    "        grads = pf.initial_velocity_derivative(ctx.solver)\n",
+    "        flat_grad = torch.zeros((len(ctx.bids), len(grads[ctx.bids[0]])), dtype=float)\n",
+    "        for id, g in grads.items():\n",
+    "            flat_grad[ctx.bids.index(id), :] = torch.tensor(g)\n",
+    "        return None, None, flat_grad"
    ]
   },
   {
@@ -68,16 +72,12 @@
     "solver1.set_log_level(2)\n",
     "solver1.load_mesh_from_settings()\n",
     "\n",
-    "mesh = solver1.mesh()\n",
-    "v = mesh.vertices()\n",
-    "vertices = torch.tensor(solver1.mesh().vertices(), requires_grad=True)\n",
-    "\n",
     "# Simulation 2\n",
     "\n",
     "config[\"initial_conditions\"][\"velocity\"][0][\"value\"] = [3, 0]\n",
     "solver2 = pf.Solver()\n",
     "solver2.set_settings(json.dumps(config), False)\n",
-    "solver2.set_log_level(2)\n",
+    "solver2.set_log_level(1)\n",
     "solver2.load_mesh_from_settings()"
    ]
   },
@@ -89,16 +89,25 @@
    "source": [
     "\n",
     "# Verify gradient\n",
-    "\n",
-    "def loss(vertices):\n",
-    "    solutions1 = Simulate.apply(solver1, vertices)\n",
-    "    solutions2 = Simulate.apply(solver2, vertices)\n",
-    "    print(obj.value(vertices.reshape(-1,1).detach().numpy()))\n",
+    "def loss(param):\n",
+    "    solutions1 = Simulate.apply(solver1, body_ids, param)\n",
+    "    solutions2 = Simulate.apply(solver2, body_ids, param)\n",
     "    return torch.linalg.norm(solutions1[:, -1]) * torch.linalg.norm(solutions2[:, -1])\n",
     "\n",
     "torch.set_printoptions(12)\n",
     "\n",
-    "param = vertices.clone().detach().requires_grad_(True)\n",
+    "dt = 0.04\n",
+    "solver1.set_cache_level(pf.CacheLevel.Derivatives)\n",
+    "solver1.build_basis()\n",
+    "solver1.assemble()\n",
+    "solver1.init_timestepping(0, dt)\n",
+    "solver2.set_cache_level(pf.CacheLevel.Derivatives)\n",
+    "solver2.build_basis()\n",
+    "solver2.assemble()\n",
+    "solver2.init_timestepping(0, dt)\n",
+    "param = torch.tensor([[5., 0], [0, 0]], requires_grad=True)\n",
+    "body_ids = [1, 3]\n",
+    "\n",
     "theta = torch.randn_like(param)\n",
     "l = loss(param)\n",
     "l.backward()\n",
@@ -109,7 +118,7 @@
     "    f1 = loss(param + theta * t)\n",
     "    f2 = loss(param - theta * t)\n",
     "    fd = (f1 - f2) / (2 * t)\n",
-    "    print(f'grad {analytic}, fd {fd} {(f1 - l) / t} {(l - f2) / t}, relative err {abs(analytic - fd) / abs(analytic):.3e}')\n",
+    "    print(f'\\ngrad {analytic}, fd {fd} {(f1 - l) / t} {(l - f2) / t}, relative err {abs(analytic - fd) / abs(analytic):.3e}')\n",
     "    print(f'f(x+dx)={f1}, f(x)={l.detach()}, f(x-dx)={f2}')\n",
     "    assert(abs(analytic - fd) <= 1e-4 * abs(analytic))"
    ]