{ "cells": [ { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "data": { "image/png": "iVBORw0KGgoAAAANSUhEUgAAAaQAAAE/CAYAAAAJ7be1AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy8QZhcZAAALwklEQVR4nO3cf8jud13H8dfbnVmhC//YDdq20xEa0hBJuBlEUZGrlojLQHBECAYH/7AmFMwaJBX+MQQJVlAHNipYSjCHghO3kbGEZrs3lm2eTYYwdlLaMREd/iGnvfvj3Mqw4znbub73ud73dT0ecMP16/5831/um+vJ93t9uaq7AwDr9qp1DwAAiSABMIQgATCCIAEwgiABMIIgATDCkXVs9Morr+xjx46tY9MArNGjjz76je7eOddzawnSsWPHsre3t45NA7BGVfXsj3rOKTsARhAkAEYQJABGECQARhAkAEYQJABGECQARhAkAEYQJABGECQARhAkAEYQJABGECQARhAkAEYQJABGECQARhAkAEYQJABGWDlIVXVNVX2+qk5W1ZNVdcsSgwGwXY4ssMaZJH/Y3Y9V1RVJHq2qB7r7ywusDcCWWPkIqbu/3t2P7d/+TpKTSa5adV0AtssSR0g/UFXHkrw1yReXXBc2xe/v3vqD23fs3b7GSWCexS5qqKrXJrknyQe7+9vneP54Ve1V1d7p06eX2iwAG2KRIFXV5Tkbo7u7+5Pnek13n+ju3e7e3dnZWWKzAGyQJa6yqyR3JjnZ3R9bfSQAttESR0i/kOR3k/xqVT2+//P2BdYFYIusfFFDd38hSS0wCwBbzDc1ADCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMIIgATCCIAEwgiABMMIiQaqqu6rq+ap6Yon1ANg+Sx0h/V2SGxdaC4AttEiQuvuhJN9cYi0AtpPPkAAY4ZIFqaqOV9VeVe2dPn36Um0WgEPikgWpu09092537+7s7FyqzQJwSDhlB8AIS132/fEk/5bkTVV1qqp+b4l1AdgeR5ZYpLtvXmIdALaXU3YAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMsEiQqurGqnq6qp6pqg8tsSYA22XlIFXVZUn+OslvJrkuyc1Vdd2q6wKwXZY4Qro+yTPd/dXu/l6STyS5aYF1AdgiSwTpqiTPveT+qf3HAOBlWyJIdY7H+v+9qOp4Ve1V1d7p06cX2CwAm2SJIJ1Kcs1L7l+d5Gs//KLuPtHdu929u7Ozs8BmAdgkSwTpkSTXVtUbq+rVSd6T5NMLrAvAFjmy6gLdfaaqPpDkc0kuS3JXdz+58mQAbJWVg5Qk3X1fkvuWWAuA7eSbGgAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhhpSBV1bur6smqerGqdpcaCoDts+oR0hNJfjvJQwvMAsAWO7LKL3f3ySSpqmWmAWBr+QwJgBEueIRUVQ8mef05nrqtuz/1cjdUVceTHE+So0ePvuwBAdgOFwxSd9+wxIa6+0SSE0myu7vbS6wJwOZwyg6AEVa97PtdVXUqyc8n+UxVfW6ZsQDYNqteZXdvknsXmgWALeaUHQAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACMIEgAjCBIAIwgSACOsFKSq+mhVPVVVX6qqe6vqdUsNBsB2WfUI6YEkb+7utyT5SpI/Xn0kALbRSkHq7vu7+8z+3YeTXL36SABsoyU/Q3pfks8uuB4AW+TIhV5QVQ8mef05nrqtuz+1/5rbkpxJcvd51jme5HiSHD169KKGBWBzXTBI3X3D+Z6vqvcmeUeSt3V3n2edE0lOJMnu7u6PfB0A2+mCQTqfqroxya1Jfrm7v7vMSLC57ti7fd0jwFirfob0V0muSPJAVT1eVX+zwEwAbKGVjpC6+2eWGgSA7eabGgAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABhBkAAYobr70m+06nSSZ1/Br1yZ5BsHNM5E9nez2d/NZn/P76e7e+dcT6wlSK9UVe119+6657hU7O9ms7+bzf5ePKfsABhBkAAY4bAE6cS6B7jE7O9ms7+bzf5epEPxGRIAm++wHCEBsOEOTZCq6i+q6ktV9XhV3V9VP7XumQ5SVX20qp7a3+d7q+p1657pIFXVu6vqyap6sao28gqlqrqxqp6uqmeq6kPrnuegVdVdVfV8VT2x7lkuhaq6pqo+X1Un9/+Xb1n3TAepqn68qv69qv5jf3//bOU1D8spu6r6ye7+9v7tP0hyXXe/f81jHZiq+vUk/9zdZ6rq9iTp7lvXPNaBqaqfTfJikr9N8kfdvbfmkRZVVZcl+UqSX0tyKskjSW7u7i+vdbADVFW/lOSFJP/Q3W9e9zwHrarekOQN3f1YVV2R5NEkv7Wpf+OqqiSv6e4XquryJF9Ickt3P3yxax6aI6Tvx2jfa5IcjpJepO6+v7vP7N99OMnV65znoHX3ye5+et1zHKDrkzzT3V/t7u8l+USSm9Y804Hq7oeSfHPdc1wq3f317n5s//Z3kpxMctV6pzo4fdYL+3cv3/9Z6X350AQpSarqI1X1XJLfSfKn657nEnpfks+uewhWclWS515y/1Q2+M1q21XVsSRvTfLF9U5ysKrqsqp6PMnzSR7o7pX2d1SQqurBqnriHD83JUl339bd1yS5O8kH1jvt6i60v/uvuS3JmZzd50Pt5ezvBqtzPLbRR/nbqqpem+SeJB/8oTM7G6e7/7e7fy5nz+BcX1UrnZo9ssxYy+juG17mS/8xyWeSfPgAxzlwF9rfqnpvknckeVsflg/7zuMV/H030akk17zk/tVJvramWTgg+5+l3JPk7u7+5LrnuVS6+1tV9S9Jbkxy0RexjDpCOp+quvYld9+Z5Kl1zXIpVNWNSW5N8s7u/u6652FljyS5tqreWFWvTvKeJJ9e80wsaP9D/juTnOzuj617noNWVTvfv/q3qn4iyQ1Z8X35MF1ld0+SN+XslVjPJnl/d//Xeqc6OFX1TJIfS/I/+w89vOFXFb4ryR1JdpJ8K8nj3f0b651qWVX19iR/meSyJHd190fWPNKBqqqPJ/mVnP026P9O8uHuvnOtQx2gqvrFJP+a5D9z9n0qSf6ku+9b31QHp6rekuTvc/b/+VVJ/qm7/3ylNQ9LkADYbIfmlB0Am02QABhBkAAYQZAAGEGQABhBkAAYQZAAGEGQABjh/wDcxUF7NcH4kgAAAABJRU5ErkJggg==\n", "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "import numpy as np\n", "from scipy.spatial.distance import pdist, squareform\n", "\n", "import matplotlib.pyplot as plt\n", "import scipy.integrate as integrate\n", "import matplotlib.animation as animation\n", "\n", "class ParticleBox:\n", " \"\"\"Orbits class\n", "\n", " init_state is an [N x 4] array, where N is the number of particles:\n", " [[x1, y1, vx1, vy1],\n", " [x2, y2, vx2, vy2],\n", " ... ]\n", "\n", " bounds is the size of the box: [xmin, xmax, ymin, ymax]\n", " \"\"\"\n", " def __init__(self,\n", " init_state = [[1, 0, 0, -1],\n", " [-0.5, 0.5, 0.5, 0.5],\n", " [-0.5, -0.5, -0.5, 0.5]],\n", " bounds = [-2, 2, -2, 2,0],\n", " size = 0.04,\n", " M = 0.05,\n", " G = 9.8):\n", " self.init_state = np.asarray(init_state, dtype=float)\n", " self.M = M * np.ones(self.init_state.shape[0])\n", " self.size = size\n", " self.state = self.init_state.copy()\n", " self.time_elapsed = 0\n", " self.bounds = bounds\n", " self.G = G\n", "\n", " def step(self, dt):\n", " \"\"\"step once by dt seconds\"\"\"\n", " self.time_elapsed += dt\n", "\n", " # update positions\n", " self.state[:, :2] += dt * self.state[:, 2:]\n", " \n", "\n", " # find pairs of particles undergoing a collision\n", " D = squareform(pdist(self.state[:, :2]))\n", " ind1, ind2 = np.where(D < 2 * self.size)\n", " unique = (ind1 < ind2)\n", " ind1 = ind1[unique]\n", " ind2 = ind2[unique]\n", "\n", " # update velocities of colliding pairs\n", " for i1, i2 in zip(ind1, ind2):\n", " # mass\n", " m1 = self.M[i1]\n", " m2 = self.M[i2]\n", "\n", " # location vector\n", " r1 = self.state[i1, :2]\n", " r2 = self.state[i2, :2]\n", "\n", " # velocity vector\n", " v1 = self.state[i1, 2:]\n", " v2 = self.state[i2, 2:]\n", "\n", " # relative location & velocity vectors\n", " r_rel = r1 - r2\n", " v_rel = v1 - v2\n", "\n", " # momentum vector of the center of mass\n", " v_cm = (m1 * v1 + m2 * v2) / (m1 + m2)\n", "\n", " # collisions of spheres reflect v_rel over r_rel\n", " rr_rel = np.dot(r_rel, r_rel)\n", " vr_rel = np.dot(v_rel, r_rel)\n", " v_rel = 2 * r_rel * vr_rel / rr_rel - v_rel\n", "\n", " # assign new velocities\n", " self.state[i1, 2:] = v_cm + v_rel * m2 / (m1 + m2)\n", " self.state[i2, 2:] = v_cm - v_rel * m1 / (m1 + m2)\n", "\n", " # check for crossing boundary\n", " crossed_x1 = (self.state[:, 0] < self.bounds[0] + self.size)\n", " crossed_x2 = (self.state[:, 0] > self.bounds[1] - self.size)\n", " crossed_y1 = (self.state[:, 1] < self.bounds[2] + self.size)\n", " crossed_y2 = (self.state[:, 1] > self.bounds[3] - self.size)\n", "\n", " crossed_d1 = (self.state[:, 0] < self.bounds[4] + self.size)&(self.state[:, 0] > self.bounds[4] - self.size) & (self.state[:, 2] < 0)\n", " crossed_d2 = (self.state[:, 0] > self.bounds[4] - self.size)&(self.state[:, 0] < self.bounds[4] + self.size)&(self.state[:, 2] > 0)\n", "\n", " self.state[crossed_x1, 0] = self.bounds[0] + self.size\n", " self.state[crossed_x2, 0] = self.bounds[1] - self.size\n", "\n", " self.state[crossed_d1, 0] = self.bounds[4] + self.size\n", " self.state[crossed_d2, 0] = self.bounds[4] - self.size\n", "\n", " self.state[crossed_y1, 1] = self.bounds[2] + self.size\n", " self.state[crossed_y2, 1] = self.bounds[3] - self.size\n", "\n", " self.state[crossed_x1 | crossed_x2, 2] *= -1\n", " self.state[crossed_y1 | crossed_y2, 3] *= -1\n", " self.state[crossed_d1 | crossed_d2, 2] *= -1\n", "\n", "\n", " # check at demon gate\n", "\n", "\n", " # add gravity\n", " # self.state[:, 3] -= self.M * self.G * dt\n", "\n", "\n", "#------------------------------------------------------------\n", "# set up initial state\n", "np.random.seed(0)\n", "init_state = -0.5 + np.random.random((50, 4))\n", "init_state[:, :2] *= 3.9\n", "\n", "box = ParticleBox(init_state, size=0.04)\n", "dt = 1. / 30 # 30fps\n", "\n", "\n", "#------------------------------------------------------------\n", "# set up figure and animation\n", "fig = plt.figure()\n", "fig.subplots_adjust(left=0, right=1, bottom=0, top=1)\n", "ax = fig.add_subplot(111, aspect='equal', autoscale_on=False,\n", " xlim=(-3.2, 3.2), ylim=(-2.4, 2.4))\n", "\n", "# particles holds the locations of the particles\n", "particles, = ax.plot([], [], 'bo', ms=6)\n", "\n", "# rect is the box edge\n", "rect = plt.Rectangle(box.bounds[::2],\n", " box.bounds[1] - box.bounds[0],\n", " box.bounds[3] - box.bounds[2],\n", " ec='none', lw=2, fc='none')\n", "ax.add_patch(rect)\n", "\n", "# display demon line\n", "bottom = [0,0]\n", "top = [-2,2]\n", "plt.plot(bottom, top, color=\"#6c3376\", linewidth=3)\n", "\n", "\n", "\n", "\n", "def init():\n", " \"\"\"initialize animation\"\"\"\n", " global box, rect\n", " particles.set_data([], [])\n", " rect.set_edgecolor('none')\n", " return particles, rect\n", "\n", "def animate(i):\n", " \"\"\"perform animation step\"\"\"\n", " global box, rect, dt, ax, fig\n", " box.step(dt)\n", "\n", " ms = int(fig.dpi * 2 * box.size * fig.get_figwidth()\n", " / np.diff(ax.get_xbound())[0])\n", "\n", " # update pieces of the animation\n", " rect.set_edgecolor('k')\n", " particles.set_data(box.state[:, 0], box.state[:, 1])\n", " particles.set_markersize(ms)\n", " return particles, rect\n", "\n", "ani = animation.FuncAnimation(fig, animate, frames=600,\n", " interval=10, blit=True, init_func=init)\n", "\n", "\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "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.7.4" } }, "nbformat": 4, "nbformat_minor": 2 }