fix memory issue, support multiple states

Author: Huangzizhou

src/state/state.cpp

@@ -370,8 +370,8 @@ void define_solver(py::module_ &m)
           py::arg("log_level"))
 
       .def(
-          "mesh", [](State &s) -> mesh::Mesh * { return s.mesh.get(); },
-          "Get mesh in simulator")
+          "mesh", [](State &s) -> mesh::Mesh& { return *s.mesh.get(); },
+          "Get mesh in simulator", py::return_value_policy::reference)
 
       .def(
           "load_mesh_from_settings",

test/test_differentiable.ipynb

@@ -20,26 +20,36 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "\n",
     "# Differentiable simulator that computes shape derivatives\n",
     "class Simulate(torch.autograd.Function):\n",
     "\n",
     "    @staticmethod\n",
-    "    def forward(ctx, solver, vertices):\n",
-    "        solver.mesh().set_vertices(vertices)\n",
-    "        solver.set_cache_level(pf.CacheLevel.Derivatives) # enable backward derivatives\n",
-    "        solver.solve()\n",
-    "        cache = solver.get_solution_cache()\n",
-    "        solutions = torch.zeros((solver.ndof(), cache.size()))\n",
-    "        for t in range(cache.size()):\n",
-    "            solutions[:, t] = torch.tensor(cache.solution(t))\n",
+    "    def forward(ctx, solvers, vertices):\n",
+    "        solutions = []\n",
+    "        for solver in solvers:\n",
+    "            solver.mesh().set_vertices(vertices)\n",
+    "            solver.set_cache_level(pf.CacheLevel.Derivatives) # enable backward derivatives\n",
+    "            solver.solve()\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",
+    "            solutions.append(sol)\n",
     "        \n",
-    "        ctx.solver = solver\n",
-    "        return solutions\n",
+    "        ctx.save_for_backward(vertices)\n",
+    "        ctx.solvers = solvers\n",
+    "        return tuple(solutions)\n",
     "\n",
     "    @staticmethod\n",
-    "    def backward(ctx, grad_output):\n",
-    "        ctx.solver.solve_adjoint(grad_output)\n",
-    "        return None, torch.tensor(pf.shape_derivative(ctx.solver))"
+    "    @torch.autograd.function.once_differentiable\n",
+    "    def backward(ctx, *grad_output):\n",
+    "        vertices, = ctx.saved_tensors\n",
+    "        grad = torch.zeros_like(vertices)\n",
+    "        for i, solver in enumerate(ctx.solvers):\n",
+    "            solver.solve_adjoint(grad_output[i])\n",
+    "            grad += torch.tensor(pf.shape_derivative(solver))\n",
+    "        return None, grad\n"
    ]
   },
   {
@@ -48,20 +58,33 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "\n",
     "root = \"../data/differentiable/input\"\n",
     "with open(root + \"/initial-contact.json\", \"r\") as f:\n",
     "    config = json.load(f)\n",
     "\n",
     "config[\"root_path\"] = root + \"/initial-contact.json\"\n",
     "\n",
-    "solver = pf.Solver()\n",
-    "solver.set_settings(json.dumps(config), False)\n",
-    "solver.set_log_level(2)\n",
-    "solver.load_mesh_from_settings()\n",
+    "# Simulation 1\n",
+    "\n",
+    "solver1 = pf.Solver()\n",
+    "solver1.set_settings(json.dumps(config), False)\n",
+    "solver1.set_log_level(2)\n",
+    "solver1.load_mesh_from_settings()\n",
+    "# solver1.solve()\n",
     "\n",
-    "mesh = solver.mesh()\n",
+    "mesh = solver1.mesh()\n",
     "v = mesh.vertices()\n",
-    "vertices = torch.tensor(solver.mesh().vertices(), requires_grad=True)"
+    "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.load_mesh_from_settings()\n",
+    "# solver2.solve()"
    ]
   },
   {
@@ -70,11 +93,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "\n",
     "# Verify gradient\n",
     "\n",
     "def loss(vertices):\n",
-    "    solutions = Simulate.apply(solver, vertices)\n",
-    "    return torch.linalg.norm(solutions[:, -1])\n",
+    "    solutions1, solutions2 = Simulate.apply([solver1, solver2], vertices)\n",
+    "    return torch.linalg.norm(solutions1[:, -1]) * torch.linalg.norm(solutions2[:, -1])\n",
     "\n",
     "torch.set_printoptions(12)\n",
     "\n",