NewtonGravitationalForce can have a line connecting the bodies as mobject

Author: jCodingStuff

example_scenes.py

@@ -743,15 +743,29 @@ def construct(self) -> None:
                 tracer=tracing_lines[i]
             ) for i in range(n_masses)
         ]
-        # forces: list[Force] = []
-        # for i, body1 in enumerate(bodies):
-        #     for _, body2 in enumerate(bodies[i+1:],start=i+1):
-        #         forces.append(
-        #             NewtonGravitationalForce((body1, body2))
-        #         )
-        # system: PhysicalSystem = PhysicalSystem(bodies, forces)
-        system: GravitationalSystem = GravitationalSystem(bodies)
-        system.fill_forces()
+        force_colors: list[str] = [YELLOW, PURPLE, LIGHT_BROWN]
+        k: int = 0
+        forces: list[Force] = []
+        for i, body1 in enumerate(bodies):
+            for _, body2 in enumerate(bodies[i+1:],start=i+1):
+                line = Line(
+                            body1.position,
+                            body2.position,
+                            stroke_color=force_colors[k],
+                            stroke_width=5
+                       )
+                self.add(line)
+                forces.append(
+                    NewtonGravitationalForce(
+                        (body1, body2),
+                        mobjects=(line,)
+                    )
+                )
+                k += 1
+        system: PhysicalSystem = PhysicalSystem(bodies, forces)
+        # system: GravitationalSystem = GravitationalSystem(bodies)
+        # system.fill_forces()
+        self.bring_to_front(*mass_circles)  # put them on top of the force mobjects
 
         # Wait a couple seconds
         self.wait(2)

manimlib/animation/physics.py

@@ -107,10 +107,13 @@ def interpolate_mobject(self, alpha: float) -> None:
             body.set_position(pos, update_mobject_position=True)
             body.update_tracer()
             if body.tracer is not None:
-                self.scene.bring_to_front(body.tracer)
+                self.scene.bring_to_back(body.tracer)
             if body.mobj is not None:
                 self.scene.bring_to_front(body.mobj)
-        # # Handle background and foreground mobjects
+        # Update mobjects in forces
+        for force in self.mobject.forces:
+            force.update_mobjects()
+        # Handle background and foreground mobjects
         if self.background_mobjects:
             self.scene.bring_to_back(*self.background_mobjects)
         if self.foreground_mobjects:

manimlib/physics/force.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from typing import Union, Callable
+from manimlib.mobject.mobject import Mobject
 
 from manimlib.physics.body import Body
 from manimlib.mobject.geometry import Line
@@ -20,26 +21,35 @@ class Force(metaclass=ABCMeta):
     def __init__(
         self,
         bodies: tuple[Body, ...],
-        line: Union[Line, Line3D]=None
+        mobjects: tuple[Mobject, ...]=()
     ) -> None:
         """
         Initialize a new Force object
 
         Keyword arguments
         -----------------
-        bodies (list[Body]): the bodies to which the force applies
-        line: (Line | Line3D): line representing the force and follows
-              the bodies, must already be set to have the bodies at its
-              extremes (default None)
+        bodies (tuple[Body, ...]): the bodies to which the force applies
+        mobjects: tuple[Mobject, ...]: mobject(s) representing the force, should
+              be already set to a desired position (subclasses know
+              how to update them). Regular shapes should be used for 2D
+              simulations and 3D shapes for 3D simulations! Otherwise
+              things MAY BREAK! (default: empty tuple)
         """
-        self.bodies: tuple[Body] = bodies
+        self.bodies: tuple[Body, ...] = bodies
         if not self.bodies:
             raise Exception("No bodies have been provided!")
-        self.line = line
+        self.mobjects: tuple[Mobject, ...] = mobjects
 
     def __str__(self) -> str:
         body_info: str = ','.join([f"body{i}_index={body.index}" for i, body in enumerate(self.bodies, start=1)])
-        return (f"{body_info},line={self.line}")
+        return (f"{body_info},mobjects={self.mobjects}")
+
+    def update_mobjects(self) -> None:
+        """
+        Update the foce mobject(s)
+        (to be implemented by each subclass)
+        """
+        pass
 
     @abstractmethod
     def apply(self, forces: np.ndarray) -> None:
@@ -63,30 +73,55 @@ class NewtonGravitationalForce(Force):
 
     def __init__(
         self,
-        bodies: tuple[Body],
-        line: Union[Line, Line3D]=None,
+        bodies: tuple[Body, Body],
+        mobjects: tuple[Union[Line, Line3D]]=(),
         G: float=1.0
     ) -> None:
         """
         Initialize a new NewtonGravitationalForce instance
 
         Keyword arguments
         -----------------
-        bodies (list[Body]): the bodies to which the force applies
-        line (Line | Line3D): line representing the force and follows
+        bodies (tuple[Body, Body]): the bodies to which the force applies
+        mobjects (tuple[Line | Line3D]): line representing the force and follows
              the bodies, must already be set to have the bodies at its
-             extremes (default None)
+             extremes (start=bodies[0], end=bodies[1]) (default: empty tuple)
         G (float): gravitational constant (default: 1.0)
         """
         if len(bodies) != 2:
             raise Exception("You must provide 2 bodies!")
-        super().__init__(bodies)
+        super().__init__(bodies, mobjects)
         self.G: float = G
 
     def __str__(self) -> str:
         return (f"{self.__class__.__name__}<{super().__str__()},"
                 f"G={self.G}>")
 
+    def update_mobjects(self) -> None:
+        """
+        Update the line mobject to follow the bodies
+        """
+        if not self.mobjects or self.mobjects[0] is None:
+            return
+        line: Union[Line, Line3D] = self.mobjects[0]
+        body1, body2 = self.bodies[0], self.bodies[1]
+        delta_pos: np.ndarray =  body2.position - body1.position
+        distance: float = np.linalg.norm(delta_pos)
+        if isinstance(line, Line):  # for 2D simulations
+            line.set_length(distance)
+            line.set_angle(np.arctan2(delta_pos[1], delta_pos[0]))
+            line.move_to((body1.position+body2.position)/2)
+        elif isinstance(line, Line3D):  # TODO: implement this
+            print(
+                f"({self.__class__.__name__}) WARNING: updating "
+                "is not implemented yet."
+            )
+        else:
+            print(
+                f"({self.__class__.__name__}) WARNING: the mobject "
+                "is not a line. Cannot update it."
+            )
+
     def apply(self, forces: np.ndarray) -> None:
         body1, body2 = self.bodies[0], self.bodies[1]
         delta_pos: np.ndarray = body2.position - body1.position

manimlib/physics/physical_system.py

@@ -29,13 +29,16 @@ def __init__(self, bodies: list[Body]=[], forces: list[Force]=[], **kwargs):
         forces (list[Force]): forces to be added to the system (default: empty list)
         """
         # Aggregate all submobjects and pass them to the superclass constructor
-        mobjects: list[Mobject] = [
-            body.mobj for body in bodies if body.mobj is not None
-        ] + [
-            body.tracer for body in bodies if body.tracer is not None
-        ] + [
-            force.line for force in forces if force.line is not None
-        ]
+        mobjects: list[Mobject] = []
+        for body in bodies:
+            if body.mobj is not None:
+                mobjects.append(body.mobj)
+            if body.tracer is not None:
+                mobjects.append(body.tracer)
+        for force in forces:
+            for mobj in force.mobjects:
+                if mobj is not None:
+                    mobjects.append(mobj)
         super().__init__(*mobjects, **kwargs)
         self.bodies: list[Body] = bodies
         for i, body in enumerate(self.bodies):  # set indices for the bodies
@@ -136,6 +139,8 @@ def fill_forces(self, **kwargs) -> None:
         """
         Add a NewtonGravitationalForce to every pair of bodies
         TODO: add an option in 'kwargs' to automatically add line mobjects for the force
+              (they will have to be added with the self.add() method as well, since this is
+              a mobject itself too)
 
         Keyword arguments
         -----------------