Simulation of a set of bodies subjected to gravity

Simulation of a set of bodies subjected to gravity#

Scope#

This notebook uses the (Point Mass Dynamics) PMD class to simulate gravitational interaction between massive objects.

Required files

Before using this notebook, download the module PMD.py and put it in your working directory.

Coding#

%matplotlib widget
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from scipy import integrate, optimize, spatial
from matplotlib import animation, rc
from PMD import PMD, distances, MetaForce

rc("animation", html="html5")
np.random.seed(333)

class Gravity(PMD):
    def __init__(self, G=6.67e-11, **kwargs):
        """
        2D gravity
        """
        self.G = G
        super().__init__(**kwargs)

    def derivative(self, X, t, cutoff_radius=1.0e-2):
        """
        Acceleration de chaque masse !
        """
        m, G = self.m, self.G
        n = len(m)
        P = X[: 2 * n].reshape(n, 2)
        V = X[2 * n :].reshape(n, 2)
        M = m * m[:, np.newaxis]
        D, R, U = distances(P)
        np.fill_diagonal(R, np.inf)
        if cutoff_radius > 0.0:
            R = np.where(R > cutoff_radius, R, cutoff_radius)
        F = ((G * M * R**-2)[:, :, np.newaxis] * U).sum(axis=0)
        A = (F.T / m).T
        X2 = X.copy()
        X2[: 2 * n] = V.flatten()
        X2[2 * n :] = A.flatten()
        return X2

Animation#

# SETUP
G = 1.0e03
nr = 3
nt = 1
nm = nr * nt + 1
m = np.ones(nm) * 4.0e-3
m[0] = 1.0
r = np.linspace(1.0, 2.0, nr)
theta = np.linspace(0.0, np.pi * 2, nt, endpoint=False)
R, Theta = np.meshgrid(r, theta)
r = np.concatenate([[0.0], R.flatten()])
theta = np.concatenate([[0.0], Theta.flatten()])

v = np.zeros_like(r)
v[1:] = (
    (G * m[0] / r[1:]) ** 0.5
    * 0.75
    * np.random.normal(loc=1.0, scale=0.05, size=nm - 1)
)
x = r * np.cos(theta)
y = r * np.sin(theta)
vx = -v * np.sin(theta)
vy = v * np.cos(theta)

P = np.array([x, y]).transpose()
V = np.array([vx, vy]).transpose()

vG = (V * m[:, np.newaxis]).sum(axis=0) / m.sum()
V -= vG

s = Gravity(m=m, P=P, V=V, G=G, nk=4000)
dt = 1.0e-3
nt = 50

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[3], line 31
     28 vG = (V * m[:, np.newaxis]).sum(axis=0) / m.sum()
     29 V -= vG
---> 31 s = Gravity(m=m, P=P, V=V, G=G, nk=4000)
     32 dt = 1.0e-3
     33 nt = 50

Cell In[2], line 7, in Gravity.__init__(self, G, **kwargs)
      3 """
      4 2D gravity
      5 """
      6 self.G = G
----> 7 super().__init__(**kwargs)

File /workspaces/positron/book/ode/ressources/PMD.py:34, in PMD.__init__(self, m, P, V, nk)
     32 self._n = n
     33 self.X = np.zeros([nk, 4 * n])
---> 34 self.X.fill(np.NAN)
     35 self.X[-1, : 2 * n] = np.array(P).flatten()
     36 self.X[-1, 2 * n :] = np.array(V).flatten()

File /opt/conda/envs/science/lib/python3.12/site-packages/numpy/__init__.py:414, in __getattr__(attr)
    411     import numpy.char as char
    412     return char.chararray
--> 414 raise AttributeError("module {!r} has no attribute "
    415                      "{!r}".format(__name__, attr))

AttributeError: module 'numpy' has no attribute 'NAN'

pcolors = "r"
tcolors = "k"


fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_aspect("equal")
margin = 1.0
plt.axis([-2, 2, -2, 2])
plt.grid()
ax.axis("off")
points = []

msize = 10.0 * (s.m / s.m.max()) ** (1.0 / 6.0)
for i in range(nm):
    plc = len(pcolors)
    pc = pcolors[i % plc]
    tlc = len(tcolors)
    tc = tcolors[i % tlc]
    (trail,) = ax.plot([], [], "-" + tc)
    (point,) = ax.plot([], [], "o" + pc, markersize=msize[i])
    points.append(point)
    points.append(trail)


def init():
    for i in range(2 * nm):
        points[i].set_data([], [])
    return points


def animate(i):
    s.solve(dt, nt)
    x, y = s.xy()
    for i in range(nm):
        points[2 * i].set_data(x[i : i + 1], y[i : i + 1])
        xt, yt = s.trail(i)
        points[2 * i + 1].set_data(xt, yt)
    return points


anim = animation.FuncAnimation(
    fig, animate, init_func=init, frames=1600, interval=20, blit=True
)


plt.close()
anim
# plt.show()
Contents