one collision per particle

guoci · Sep 19, 2023 · 5db3e7c · 5db3e7c
1 parent d058b55
commit 5db3e7c
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/Python GPU/multibody_boltzmann.ipynb b/Python GPU/multibody_boltzmann.ipynb
@@ -500,6 +500,13 @@
     "    vs[0] = v\n",
     "    for i in range(1,ts):\n",
     "        ic = id_pairs[get_deltad2_pairs(r, ids_pairs) < d_cutoff**2]\n",
+    "        mod = torch.zeros(n_particles)\n",
+    "        ic2 = []\n",
+    "        for a, b in ic:\n",
+    "            if mod[a] == mod[b] == 0:\n",
+    "                ic2.append([a, b])\n",
+    "                mod[a] = mod[b] = 1\n",
+    "        ic = torch.as_tensor(ic2)\n",
     "        v[:,ic[:,0]], v[:,ic[:,1]] = compute_new_v(v[:,ic[:,0]], v[:,ic[:,1]], r[:,ic[:,0]], r[:,ic[:,1]])\n",
     "        \n",
     "        v[0,r[0]>1] = -torch.abs(v[0,r[0]>1])\n",