sanity pass on n-body with tests

Author: parmentelat

notebooks/numpy-tps/n-body/.teacher/README-n-body-corrige-nb.md

@@ -24,11 +24,11 @@ date: 2025-09-20
 
 pour faire cette activité sur votre ordi localement, {download}`commencez par télécharger le zip<ARTEFACTS-n-body.zip>`
 
-dans ce TP on vous invite à écrire un simulateur de la trajectoire de n corps qui interagissent entre eux au travers de leur poids, pour produire des sorties de ce genre
+dans ce TP on vous invite à écrire un simulateur de la trajectoire de n corps qui interagissent entre eux au travers de leurs masses, pour produire des sorties de ce genre
 
 ```{image} media/init3.png
 :align: center
-:width: 400px
+:width: 600px
 ```
 
 +++
@@ -44,23 +44,18 @@ on suppose:
 
 ## imports
 
-importer les librairies qui vont bien pour cet exercice; on utilisera le mode `ipympl` de `matplotlib`
+on pourra utiliser le mode `ipympl` de `matplotlib`
 
 ```{code-cell} ipython3
-# à vous
-```
-
-```{code-cell} ipython3
-# prune-cell
 import numpy as np
 import matplotlib.pyplot as plt
 
 %matplotlib ipympl
 ```
 
-## initialisation
+## initialisation aléatoire
 
-en fixant arbitrairement des limites dans l'espace des positions, des vitesses et des masses, on tire au hasard une configuration de départ pour la simulation
+en fixant arbitrairement des limites dans l'espace des positions, des vitesses et des masses, la fonction `init_problem()` tire au hasard une configuration de départ pour la simulation
 
 ```{code-cell} ipython3
 # les bornes pour le tirage au sort initial
@@ -75,9 +70,11 @@ speed_max = 1.
 ```{code-cell} ipython3
 # votre code
 
-def init_problem(n):
+def init_problem(N):
     """
     retourne un tuple masses, positions, speeds
+    de formes resp.   (N,)    (2, N)     (2, N)
+    tiré au sort dans les bornes définies ci-dessus
     """
     return None, None, None
 ```
@@ -86,20 +83,19 @@ def init_problem(n):
 # prune-cell
 
 def init_problem(N):
-    """
-    retourne un tuple masses, positions, speeds
-    """
     masses = np.random.uniform(0, mass_max, size=N)
     X = np.random.uniform(x_min, x_max, N)
     Y = np.random.uniform(y_min, y_max, N)
     angles = np.random.uniform(0, 2*np.pi, N)
     modules = np.random.uniform(0, speed_max, N)
     Sx = modules * np.cos(angles)
     Sy = modules * np.sin(angles)
-    return (masses,
-            np.concatenate((X, Y)).reshape((2, -1)), 
-            np.concatenate((Sx, Sy)).reshape((2, -1))
-           )
+    return (
+        masses,
+        # reshape(2, -1) is equivalent to (2, N) here
+        np.concatenate((X, Y)).reshape((2, -1)), 
+        np.concatenate((Sx, Sy)).reshape((2, -1))
+       )
 ```
 
 ```{code-cell} ipython3
@@ -117,6 +113,30 @@ except:
     print("KO")
 ```
 
+## initialisation reproductible
+
+par commodité on vous donne la fonction suivante qui crée 3 objets:
+
+- le premier - pensez au soleil - de masse 3, an centre de la figure, de vitesse nulle
+- et deux objets de masse 1, disposés symétriquement autour du soleil  
+  - position initiale (5, 1) et vitesse initiale (-1, 0)
+  - symétrique en     (-5, -1) et vitesse initiale (1, 0)
+
+```{code-cell} ipython3
+# for your convenience
+
+def init3():
+    # first element is sun-like: heavy, at the center, and no speed
+    masses = np.array([3, 1, 1], dtype=float)
+    positions = np.array([
+        [0, 5, -5], 
+        [0, 1, -1]], dtype=float)
+    speeds = np.array([
+        [0, -1, 1], 
+        [0, 0, 0]], dtype=float)
+    return masses, positions, speeds
+```
+
 ## les forces
 
 à présent, on va écrire un fonction qui va calculer les influences de toutes les particules entre elles, suivant la loi de Newton
@@ -132,7 +152,7 @@ pour cela on se propose d'écrire la fonction suivante
 ```{code-cell} ipython3
 # votre code
 
-def forces(positions, masses, G=1.0):
+def forces(masses, positions, G=1.0):
     """
     returns an array of shape (2,  N)
     that contains the force felt by each mass from all the others
@@ -145,7 +165,7 @@ def forces(positions, masses, G=1.0):
 ```{code-cell} ipython3
 # prune-cell
 
-def forces(positions, masses, G=1.0):
+def forces(masses, positions, G=1.0):
     
     # creating an extra dimension so we can do the n-square difference
     delta = positions.T[None, :, :] - positions.T[:, None, :]  # (N, N, 2)
@@ -165,6 +185,19 @@ def forces(positions, masses, G=1.0):
     return forces.T  # shape (2, N)
 ```
 
+```{code-cell} ipython3
+# pour tester, voici les valeurs attendues avec la config prédéfinie
+
+masses, positions, speeds = init3()
+
+f = forces(masses, positions)
+
+# should be true
+np.all(np.isclose(f, np.array([
+    [ 0.        , -0.12257258,  0.12257258],
+    [ 0.        , -0.02451452,  0.02451452]])))
+```
+
 ## le simulateur
 
 à présent il nous reste à utiliser cette brique de base pour "faire avancer" le modèle depuis son état initial et sur un nombre fixe d'itérations
@@ -177,7 +210,7 @@ cela pourrait se passer dans une fonction qui ressemblerait à ceci
 def simulate(masses, positions, speeds, dt=0.1, nb_steps=100):
     """
     should return the positions across time
-    so an array of shape (nb_steps, N, 2)
+    so an array of shape (nb_steps, 2, N)
     optional dt is the time step
     """
     pass
@@ -189,34 +222,66 @@ def simulate(masses, positions, speeds, dt=0.1, nb_steps=100):
 def simulate(masses, positions, speeds, dt=0.1, nb_steps=100):
     # initialize result
     N = masses.shape[0]
-    result = np.empty((nb_steps, N, 2))
+    result = np.empty((nb_steps, 2, N))
     # store initial position
-    result[0] = positions.T
+    result[0] = positions
 
     # because we have stored the initial state
     # we need to start at 1 here, we dont have space for more
     for step_counter in range(1, nb_steps):
-        # compute forces
-        accelerations = forces(positions, masses) / masses
+        # compute forces and accelerations
+        accelerations = forces(masses, positions) / masses
+
         # apply on speeds
-        speeds = speeds * (1 + dt*accelerations)
+        speeds += dt*accelerations
+
         # apply on positions
-        positions = positions * (1 + dt*speeds)
+        positions += dt*speeds
+
         # store current position
-        result[step_counter] = positions.T
+        result[step_counter] = positions
     return result
 ```
 
+```{code-cell} ipython3
+# pour tester
+
+SMALL_STEPS = 4
+
+s = simulate(masses, positions, speeds, nb_steps=SMALL_STEPS)
+
+try:
+    if s.shape == (SMALL_STEPS, 2, 3):
+        print("shape OK")
+except Exception as exc:
+    print(f"OOPS {type(exc)} {exc}")
+```
+
+```{code-cell} ipython3
+# pour tester: should be true
+
+positions1 = s[1]
+
+np.all(np.isclose(positions1, np.array([
+    [ 0.        ,  4.89877427, -4.89877427],
+    [ 0.        ,  0.99975485, -0.99975485]
+])))
+```
+
 ## dessiner
 
-ne reste plus qu'à dessiner
+ne reste plus qu'à dessiner; quelques indices potentiels:
+
+- chaque corps a une couleur; l'appelant peut vous passer un jeu de couleurs, sinon en tirer un au hasard
+- pour l'épaisseur de chaque point, on peut imaginer utiliser la masse de l'objet
+  ou peut-être aussi, à tester, la vitesse de l'objet (plus c'est lent et plus on l'affiche en gros ?)
 
 ```{code-cell} ipython3
 # votre code
 
-def draw(simulation, colors=None):
+def draw(simulation, masses, colors=None):
     """
-    takes as input the result of simulate() above, 
+    takes as input the result of simulate() above,
     and draws the nb_steps positions of each of the N bodies
     ideally it should return a matplotlib Axes object
 
@@ -229,9 +294,9 @@ def draw(simulation, colors=None):
 ```{code-cell} ipython3
 # prune-cell
 
-def draw(simulation, colors=None):
+def draw(simulation, masses, colors=None):
 
-    N = simulation.shape[1]
+    nb_steps, _, N = simulation.shape
     # use colors if provided, else randomize
     colors = (colors if colors is not None
               # not too dark
@@ -240,78 +305,61 @@ def draw(simulation, colors=None):
     # create a figure and its axes    
     fig, ax = plt.subplots()
 
-    # stay safe
-    ax.set_aspect("equal")
-    ax.axis("off")
+    # not reaslly needed but may come in handy though
+    # ax.axis("off")
+    # ax.set_aspect('equal')
 
 
-    nb_steps = simulation.shape[0]
     ax.set_title(f"we have {N} bodies over {nb_steps} steps")
     for step in range(nb_steps):
-        ax.scatter(simulation[step, :, 0], simulation[step, :, 1], color=colors, lw=1)
+        ax.scatter(simulation[step, 0, :], simulation[step, 1, :], c=colors, s=masses)
     return ax
 ```
 
-## putting it all together
-
-on se met dans un état initial reproductible - surtout pour que vous puissiez facilement votre code
+## un jeu de couleurs
 
 ```{code-cell} ipython3
-import numpy as np
+# for convenience
 
-def init3():
-    # N = 3
-    # the first element is heavy, at the center, and has no speed
-    masses = np.array([3, 1, 1], dtype=float)
-    positions = np.array([[0, 5, -5],[0, 1, -1]], dtype=float)
-    speeds = np.array([[0, -1, 1],[0, 0, 0]], dtype=float)
-    return masses, positions, speeds
+colors3 = np.array([
+    [32, 32, 32],
+    (228, 90, 146),
+    (111, 0, 255),
+]) / 255
 ```
 
+## putting it all together
+
+pour commencer et tester, on se met dans l'état initial reproductible
+
 ```{code-cell} ipython3
 # décommentez ceci pour tester votre code
 
 # masses, positions, speeds = init3()
-# s = simulate(masses, positions, speeds)
-# print(f"{s.shape=}")
-# draw(s);
-```
-
-et avec ces données vous devriez obtenir plus ou moins une sortie de ce genre
-```{image} media/init3.png
+# draw(simulate(masses, positions, speeds), masses, colors)
 ```
 
 ```{code-cell} ipython3
 # prune-cell
 
-masses, positions, speeds = init3()
-draw(simulate(masses, positions, speeds));
-```
-
-```{code-cell} ipython3
-# prune-begin
-```
-
-```{code-cell} ipython3
-N = 3
-colors = np.random.uniform(0.3, 1., size=(N, 3))
-```
+steps = 100
 
-```{code-cell} ipython3
-fig, ax = plt.subplots()
-
-for time in range(s.shape[0]):
-    ax.scatter(s[time,:,0], s[time,:,1],color=colors, lw=1)
-```
+masses, positions, speeds = init3()
+s = simulate(masses, positions, speeds, nb_steps=steps)
+draw(s, masses, colors3);
 
-```{code-cell} ipython3
-s.shape
+plt.savefig("init3.png")
 ```
 
-```{code-cell} ipython3
-s.max()
+et avec ces données vous devriez obtenir plus ou moins une sortie de ce genre
+```{image} media/init3.png
 ```
 
 ```{code-cell} ipython3
+# après vous avez le droit de vous enhardir avec des scénarii plus compliqués
+# par exemple
 
+# m5, p5, s5 = init_problem(5)
+# sim5 = simulate(m5, p5, s5, nb_steps=100)
+# draw(sim5, m5);
 ```