Reworked the integration of a system

And made WaterMolecule2DExample a "complete" scene

Author: jCodingStuff

example_scenes.py

@@ -906,9 +906,12 @@ def construct(self) -> None:
                 [3., -1., 0.],  # position of mass1
                 [0., 2., 0.],  # position of mass2
                 [-3., -1., 0.],  # position of mass3
-                [0.5, 0., 0.],  # velocity of mass1
-                [0., 0.05, 0.],  # velocity of mass2
-                [-0.2, -0.2, 0.],  # velocity of mass3
+                # [0.5, 0., 0.],  # velocity of mass1
+                [0., 0., 0.],  # velocity of mass1
+                # [0., 0.05, 0.],  # velocity of mass2
+                [0., 0., 0.],  # velocity of mass2
+                # [-0.2, -0.2, 0.],  # velocity of mass3
+                [0., 0., 0.],  # velocity of mass3
             ]
         )
         masses: np.ndarray = np.array(
@@ -975,7 +978,7 @@ def construct(self) -> None:
                 HarmonicBondForce(
                     (body1, body2),
                     mobjects=(line,),
-                    k=0.15,
+                    k=0.3,
                     r0=2.9
                 )
             )
@@ -1014,19 +1017,26 @@ def construct(self) -> None:
                 )
             )
         # LangevinHeatBathForce(s)
-        # T: float = 0.1
-        # kb: float = 0.1
-        # global_seed: int = 1
-        # generator: np.random.Generator = np.random.default_rng(seed=global_seed)
-        # for body in bodies:
-        #     forces.append(
-        #         LangevinHeatBathForce(
-        #             (body,),
-        #             seed=generator.integers(0, 9999),
-        #             kb=kb,
-        #             T=T
-        #         )
-        #     )
+        temp: float = 20.0
+        global_seed: int = 4
+        max_magnitude: float = 60.0
+        generator: np.random.Generator = np.random.default_rng(seed=global_seed)
+        for body in bodies:
+            forces.append(
+                LangevinHeatBathForce(
+                    (body,),
+                    seed=generator.integers(0, 9999),
+                    T=temp,
+                    max_magnitude=max_magnitude
+                )
+            )
+        # LangevinFrictionForce(s)
+        for body in bodies:
+            forces.append(
+                LangevinFrictionForce(
+                    (body,)
+                )
+            )
         # Create system
         system: PhysicalSystem = PhysicalSystem(bodies, forces)
         # Make sure the system mobjects are positioned correctly

manimlib/__init__.py

@@ -53,9 +53,10 @@
 from manimlib.mobject.value_tracker import *
 from manimlib.mobject.vector_field import *
 
-from manimlib.physics.physical_system import *
-from manimlib.physics.force import *
 from manimlib.physics.body import *
+from manimlib.physics.force import *
+from manimlib.physics.integrator import *
+from manimlib.physics.physical_system import *
 
 from manimlib.scene.interactive_scene import *
 from manimlib.scene.scene import *
@@ -70,9 +71,9 @@
 from manimlib.utils.file_ops import *
 from manimlib.utils.images import *
 from manimlib.utils.iterables import *
-from manimlib.utils.math import *
 from manimlib.utils.paths import *
 from manimlib.utils.rate_functions import *
 from manimlib.utils.simple_functions import *
 from manimlib.utils.sounds import *
 from manimlib.utils.space_ops import *
+from manimlib.utils.trigonometry import *

manimlib/animation/physics.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
 import numpy as np
-from scipy.integrate import odeint
 
 from typing import TYPE_CHECKING
 
 from manimlib.mobject.mobject import Mobject
 from manimlib.scene.scene import Scene
 from manimlib.animation.animation import Animation
 from manimlib.physics.physical_system import PhysicalSystem
+from manimlib.physics.integrator import symplectic_euler
 from manimlib.constants import DIMENSIONS, DEFAULT_FPS
 
 if TYPE_CHECKING:
@@ -26,13 +26,11 @@ def __init__(
         mobject: PhysicalSystem,
         integrator: Callable[
             [
-                Callable[[np.ndarray, float, PhysicalSystem], np.ndarray],
-                np.ndarray,
-                np.ndarray,
-                tuple
+                PhysicalSystem,
+                np.ndarray
             ],
-            np.ndarray
-        ]=odeint,
+            tuple[np.ndarray, np.ndarray]
+        ]=symplectic_euler,
         t: np.ndarray=np.linspace(0, 10, DEFAULT_FPS*100),
         scene: Scene=None,
         background_mobjects: tuple[Mobject]=(),
@@ -44,26 +42,18 @@ def __init__(
 
         Keyword arguments
         -----------------
-        integrator (Callable[[], np.ndarray]): 
         mobject (PhysicalSystem): the PhysicalSystem mobject
         integrator (Callable[
                 [
-                    Callable[[np.ndarray, float, PhysicalSystem], np.ndarray],
-                    np.ndarray,
-                    np.ndarray,
-                    tuple
+                    PhysicalSystem,
+                    np.ndarray
                 ],
-                np.ndarray
-            ]): function in terms of (dydt, y0, t, args) returning a matrix where each row
-                is the state at every point in the t vector.
-                dydt (provided by this class) is a function in terms of the state y0 at a
-                certain time and the physical system (coming from args[0]), returning the
-                derivative (vector) of the state with respect to time.
-                y0 is the initial state (vector).
-                t is the vector of time-points in which to get the state.
-                args is a tuple containing any extra arguments dydt may need.
-                For an example of the 'integrator' function, see odeint in scipy.integrate, which
-                is the default value.
+                tuple[np.ndarray, np.ndarray]
+            ]): function in terms of (system, t) returning two tensors (pos, vel) where
+                pos[i,j,k] represents the component along the k-th axis of the position of
+                body j at time-point i (same holds por vel).
+                For an example of the 'integrator' function, see symplectic_euler in
+                manimlib.physics.integrator, which is the default value.
         t (np.ndarray[N]): vector containing all time-points of integration, must be monotonic
                            (default np.linspace(0, 10, DEFAULT_FPS*100))
         scene (Scene): scene is needed if we want to keep the body mobject, force mobject,
@@ -89,73 +79,19 @@ def __init__(
         self.scene: Scene = scene
         self.background_mobjects: tuple[Mobject] = background_mobjects
         self.foreground_mobjects: tuple[Mobject] = foreground_mobjects
-        # Assemble initial state (positions and velocities)
-        y0: np.ndarray = np.concatenate(
-            [
-                body.position for body in mobject.bodies
-            ] + [
-                body.velocity for body in mobject.bodies
-            ]
-        )
-        # Integrate (single element tuples must have a comma at the end)
-        self.y: np.ndarray = integrator(self.dydt, y0, t, (self.mobject,))
+        # Integrate the system
+        self.positions, self.velocities = integrator(self.mobject, t)
 
     def interpolate_mobject(self, alpha: float) -> None:
         row_index: int = min(
-            int(np.floor(alpha*self.y.shape[0])),
-            self.y.shape[0]-1
+            int(np.floor(alpha*self.positions.shape[0])),
+            self.positions.shape[0]-1
         )
-        state: np.ndarray = self.y[row_index,:]
-        for i, body in enumerate(self.mobject.bodies):
-            pos = state[i*DIMENSIONS:(i+1)*DIMENSIONS]
-            vel = state[(i+self.n_bodies)*DIMENSIONS:(i+self.n_bodies+1)*DIMENSIONS]
-            body.set_velocity(vel)
-            body.set_position(pos, update_mobject_position=False)  # Will update mobject later
+        self.mobject.update_positions(self.positions[row_index])
+        self.mobject.update_velocities(self.velocities[row_index])
         # Update mobjects in the system
         self.mobject.update_mobjects(
             self.scene,
             self.background_mobjects,
             self.foreground_mobjects
         )
-
-    @staticmethod
-    def dydt(
-        y: np.ndarray,
-        t: float,
-        system: PhysicalSystem
-    ) -> np.ndarray:
-        """
-        Compute the derivative of the state a physical system
-
-        Keyword arguments
-        -----------------
-        y (np.ndarray[nbodies*2*DIMENSIONS]): state containing
-          the position of each body and then its velocity
-        t (float): current time
-        system (PhysicalSystem): the physical system instance
-
-        Returns
-        -----------------
-        The derivative of the state vector
-        (np.ndarray[nbodies*2*DIMENSIONS]) containing the
-        velocities and then accelerations
-        """
-        n_bodies: int = system.get_n_bodies()
-        masses: np.ndarray = system.get_masses()
-        # Split info from state vector
-        velocities: np.ndarray = y[DIMENSIONS*n_bodies:].reshape((n_bodies, DIMENSIONS))
-        # Reshape positions so that every row contains the location of that mass
-        positions: np.ndarray = y[:DIMENSIONS*n_bodies].reshape((n_bodies, DIMENSIONS))
-        # Update the positions and velocities of the bodies in the system
-        system.update_positions(positions)
-        system.update_velocities(velocities)
-        # Create structure for forces
-        forces: np.ndarray = np.zeros((n_bodies, DIMENSIONS))
-        # Apply the forces in the system
-        for force in system.forces:
-            force.apply(forces)
-        print(f"{forces}\n")
-        # Get the accelerations (a=F/m)
-        accelerations: np.ndarray = forces/masses.reshape((n_bodies, 1))
-        # Return derivative of the state
-        return np.concatenate((velocities.flatten(), accelerations.flatten()))

manimlib/physics/force.py

@@ -7,7 +7,8 @@
 from typing import Union
 from manimlib.constants import DIMENSIONS, KB, PI, ROOM_TEMPERATURE, TAU, EPS0, NO_SEED
 from manimlib.mobject.mobject import Mobject
-from manimlib.utils.math import arctan2
+from manimlib.utils.simple_functions import cap_magnitude
+from manimlib.utils.trigonometry import arctan2
 
 from manimlib.physics.body import Body
 from manimlib.mobject.geometry import Line, Arc
@@ -319,8 +320,9 @@ def __init__(
         tridimensional: bool=False,
         seed: int=NO_SEED,
         gamma: float=1.0,
-        kb: float=KB,
-        T: float=ROOM_TEMPERATURE
+        kb: float=1.0,
+        T: float=1.0,
+        max_magnitude: float=None,
     ) -> None:
         """
         Initialize a new LangevinHeatBathForce object
@@ -337,14 +339,17 @@ def __init__(
         seed (int): seed for the random number generator. If -1, then no seed is
              passed to the generator, yielding non-reproducible results (default: -1)
         gamma (float): friction coefficient (default: 1.0)
-        kb (float): Boltzmann constant (default: 1.380649e-23)
-        T (float): temperature (default: 293.0)
+        kb (float): Boltzmann constant (default: 1.0)
+        T (float): temperature (default: 1.0)
+        max_magnitude (float): maximum magnitude allowed for a generated random force
+                      (default: None, disabled)
         """
         self.tridimensional: bool = tridimensional
         self.generator: np.random.Generator = np.random.default_rng(seed=None if seed==NO_SEED else seed)
         self.gamma: float = gamma
         self.kb: float = kb
         self.T: float = T
+        self.max_magnitude: float = max_magnitude
         super().__init__(bodies, mobjects)
 
     def __str__(self) -> str:
@@ -360,7 +365,13 @@ def apply(self, forces: np.ndarray) -> None:
         deviation: float = np.sqrt(2*body.mass*self.gamma*self.kb*self.T)
         force: np.ndarray = self.generator.normal(scale=deviation, size=DIMENSIONS)
         if not self.tridimensional:
-            force[2] == 0.0
+            force[2] = 0.
+        if self.max_magnitude is not None and cap_magnitude(force, self.max_magnitude):
+            print(
+                f"({self.__class__.__name__}) WARNING: "
+                f"force was capped to maximum "
+                f"magnitude ({self.max_magnitude})"
+            )
         forces[body.index] += force
 
 

manimlib/physics/integrator.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import numpy as np
+from manimlib.constants import DIMENSIONS
+
+from manimlib.physics.physical_system import PhysicalSystem
+
+
+def symplectic_euler(system: PhysicalSystem, t: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Perform symplectic euler integration of the physical system
+
+    Keyword arguments
+    -----------------
+    system (PhysicalSystem): the system to integrate
+    t (np.ndarray[N]): vector containing all time-points of integration, must be monotonic
+
+    Returns
+    -----------------
+    positions (np.ndarray[N, n_bodies, DIMENSIONS]): the position of each body for every
+              time-point in t
+    velocities (np.ndarray[N, n_bodies, DIMENSIONS]): the velocity of each body for every
+              time-point in t
+    """
+    n_bodies: int = system.get_n_bodies()
+    n_tpoints: int = t.shape[0]
+    positions: np.ndarray = np.zeros((n_tpoints, n_bodies, DIMENSIONS))
+    velocities: np.ndarray = np.zeros((n_tpoints, n_bodies, DIMENSIONS))
+    # Set initial positions and velocities
+    positions[0] = system.get_positions()
+    velocities[0] = system.get_velocities()
+    # Integrate
+    for i in range(1, n_tpoints):
+        # Compute dt (delta time)
+        dt: float = t[i] - t[i-1]
+        # Compute accelerations
+        system.update_positions(positions[i-1])
+        system.update_velocities(velocities[i-1])
+        accelerations: np.ndarray = system.compute_accelerations()
+        # Record new velocities
+        velocities[i] = velocities[i-1] + dt * accelerations
+        # Record new positions
+        positions[i] = positions[i-1] + dt * velocities[i]
+    # Set the initial state back in the system
+    system.update_positions(positions[0])
+    system.update_velocities(velocities[0])
+    # Return
+    return positions, velocities