Module cosmicBallet.utils.Visualization

Classes

class Visualize (celestial_objects: list, visualization_type: str = 'scientific', save_figure: bool = False, figure_name: str = None)

A class for Visualizing the trajectory of the orbits of the N-Bodies Simulated

Attributes

celestial_objects : object
An object belonging to Simulation Classes containing the simulation results for the orbital trajectory.
visual_type : str
A user input to determine whether to render an animation, or a simple scientific plot
save_fig : bool
A user input to determine whether to save the trajectory visualization. Defaults to False
figure_name : str
The name for the figure to be saved. Defaults to None
ani_name : str
The name for the animation to be saved. Defaults to None

Constructor for the Visualization class

Args

celestial_objects : object
An object of the simulation class containing the trajectory of the bodies
visualization_type : str, optional
Type of Visualization (scientific or animation). Defaults to "scientific".
save_figure : bool, optional
A value to determine whether to save the visualization or not. Defaults to False.
figure_name : str, optional
Name of the file that will hold the visualization if save_figure is True. Defaults to None.

Raises

TypeError
When celestial_objects is not a list of objects
TypeError
When visualization_type is not a string
TypeError
When save_figure is not a boolean
TypeError
When figure_name is not a string
ValueError
When visualization_type is not scientific/animation.
Expand source code
class Visualize():
    """A class for Visualizing the trajectory of the orbits of the N-Bodies Simulated

    Attributes:
        celestial_objects (object): An object belonging to Simulation Classes containing the simulation results for 
                                    the orbital trajectory.
        visual_type (str): A user input to determine whether to render an animation, or a simple scientific plot
        save_fig (bool): A user input to determine whether to save the trajectory visualization. Defaults to False
        figure_name (str): The name for the figure to be saved. Defaults to None
        ani_name (str): The name for the animation to be saved. Defaults to None
    """
    def __init__(self, celestial_objects:list, visualization_type:str="scientific", save_figure:bool=False,
                 figure_name:str=None) -> None:
        """Constructor for the Visualization class

        Args:
            celestial_objects (object): An object of the simulation class containing the trajectory of the bodies
            visualization_type (str, optional): Type of Visualization (scientific or animation). Defaults to "scientific".
            save_figure (bool, optional): A value to determine whether to save the visualization or not. Defaults to False.
            figure_name (str, optional): Name of the file that will hold the visualization if save_figure is True. Defaults to None.

        Raises:
            TypeError: When celestial_objects is not a list of objects
            TypeError: When visualization_type is not a string
            TypeError: When save_figure is not a boolean
            TypeError: When figure_name is not a string
            ValueError: When visualization_type is not scientific/animation.
        """
        try:
            assert isinstance(celestial_objects, object)
        except AssertionError:
            raise TypeError("celestial_objects needs to be a list of objects for visualization")
        try:
            assert isinstance(visualization_type, str)
        except AssertionError:
            raise TypeError("visualization_type needs to be a string")
        try:
            assert isinstance(save_figure, bool)
        except AssertionError:
            raise TypeError("The attribute save_figure must be a boolean")
        if figure_name is not None:
            try:
                assert isinstance(figure_name, str)
            except AssertionError:
                raise TypeError("The figure_name must be a string")
        try:
            accepted_choices = ["scientific", "animation"]
            assert (visualization_type.lower() in accepted_choices)
        except AssertionError:
            raise ValueError("The visualization type must be either scientific/animation")
        self.celestial_objects = celestial_objects
        self.visual_type = visualization_type
        self.save_fig = save_figure
        self.ani_name = None
        self.n_total = int(0.1*len(self.celestial_objects[0].trajectory))
        if figure_name is not None and visualization_type == "scientific":
            self.figure_name = figure_name + ".jpeg"
            self.ani_name = figure_name + ".mp4"
        else:
            if figure_name is not None:
                self.ani_name = figure_name + ".mp4"

    def _create_figure(self,celestial_objects:list, var:int, zero_padding:str)->None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
        ax.set_xlabel("x [m]")
        ax.set_ylabel("y [m]")
        ax.set_zlabel("z [m]")
        for i in range(len(celestial_objects)):
            body = celestial_objects[i]
            trajectory = np.array(body.trajectory[:int(10*var) + 1])
            ax.plot(trajectory[:, 1], trajectory[:, 2], trajectory[:, 3], label=body.name, color=body.color)
            ax.plot(trajectory[-1, 1], trajectory[-1, 2], trajectory[-1, 3], "o", color=body.color)
        ax.legend(loc='center left', bbox_to_anchor=(1.09, 0.35))
        plt.savefig(f"temp/data_{var:0{zero_padding}d}.png", dpi=300)
        #plt.show()
        plt.clf()
        plt.close(fig)

    def __create_matplotlib_animation(self):
        """Private method of the Visualization class that generates an animation of the trajectory of the bodies.       
        """
        try:
            os.mkdir("temp")
        except:
            pass
        zero_padding = len(str(self.n_total))
        for j in range(0, self.n_total):
            var = j
            self._create_figure(self.celestial_objects, var, zero_padding)
        ffmpeg_command = (
            f"ffmpeg -r {15} -i temp/data_%0{zero_padding}d.png "
            f"-vcodec mpeg4 -qscale:v 2 -filter:v 'setpts={1/1}*PTS' -y {self.ani_name}"
        )
        os.system(ffmpeg_command)
        shutil.rmtree("temp")
            
    
    def __scientific_plot(self, animate:bool)->None:
        """Private method of the Visualization class that generates a maplotlib plot of the trajectory of the bodies.

        Raises:
            ValueError: When no trajectory for the celestial objects are found
            ValueError: When no figure name is provided to save the animation
        """
        try:
            for body in self.celestial_objects:
                assert body.trajectory!=[]
        except AssertionError:
            raise ValueError(f"Empty trajectory found for {body.name}. Run Simulations first")
        if animate:
            if self.ani_name is None:
                raise ValueError("Please provide a figure name to save the animation. Animations are saved as videos" 
                                 " and require a name")
            self.__create_matplotlib_animation()
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
        i = 0
        for body in self.celestial_objects:
            trajectory = np.array(body.trajectory)
            ax.plot(trajectory[:,1], trajectory[:,2], trajectory[:,3], label=body.name, color=body.color)
            ax.scatter(trajectory[-1,1], trajectory[-1,2], trajectory[-1,3], color=body.color)
            i += 1
        plt.legend()
        ax.set_xlabel("x [m]")
        ax.set_ylabel("y [m]")
        ax.set_zlabel("z [m]")
        if self.save_fig:
            plt.savefig(self.figure_name, dpi=300)
        plt.show()
    
    def __animation(self)->None:
        print(self.celestial_objects)
        fig = mlab.figure(size=(1280, 720))
        points = []
        traj = []
        for body in self.celestial_objects:
            spacetime_point = body.trajectory[0]
            spacetime_trajectory = np.array(body.trajectory[:1])
            if spacetime_point[0] == 1:
                points.append(mlab.points3d(spacetime_point[1], spacetime_point[2], spacetime_point[3], 
                                            scale_factor=body.radius, color=body.color_myv, resolution=20))
                traj.append(mlab.plot3d(spacetime_trajectory[0,1], spacetime_trajectory[0,2], 
                                        spacetime_trajectory[0,3], color=body.color_myv))
            else:
                points.append(None)
                traj.append(None)
        
        def update(frame):
            all_points = []
            for point in points:
                if point is not None:
                    point.remove()
            for i,body in enumerate(self.celestial_objects):
                spacetime_point = body.trajectory[frame]
                if int(spacetime_point[0]) == frame+1:
                    spacetime_trajectory = np.array(body.trajectory[:frame+1])
                    points[i] = mlab.points3d(spacetime_point[1], spacetime_point[2], spacetime_point[3],
                                            scale_factor=body.radius, color=body.color_myv, resolution=20)
                    traj[i] = mlab.plot3d(spacetime_trajectory[:,1], spacetime_trajectory[:,2], 
                                           spacetime_trajectory[:,3], color=body.color_myv, tube_radius=50)
                    all_points.append(spacetime_trajectory)
            all_points = np.array(all_points)
            mlab.view(focalpoint=(all_points[:,0,1].mean(), all_points[:,0,2].mean(), all_points[:,0,3].mean()))

        if self.save_fig:
            try:
                os.mkdir("temp")
            except:
                pass
        
        zero_padding = len(str(self.n_total))

        for frame in range(self.n_total):
            update(frame*10)
            mlab.process_ui_events()
            mlab.draw()
            if self.save_fig:
                mlab.savefig(f"temp/data_{frame:0{zero_padding}d}.png")
        
        mlab.show()

        if self.save_fig:
            ffmpeg_command = (
                            f"ffmpeg -r {15} -i temp/data_%0{zero_padding}d.png "
                            f"-vcodec mpeg4 -qscale:v 2 -filter:v 'setpts={1/1}*PTS' -y {self.ani_name}"
                            )
            os.system(ffmpeg_command)
            shutil.rmtree("temp")
    
    def visualize(self, animate:bool=False)->None:
        """Method of the visualization class that generates the visualization for the trajectories based on the visualization type

        Args:
            animate (bool): Used for scientific visualization and control whether an animation is generated. Defaults to False

        Raises:
            TypeError: When the animate not a boolean value.
        """
        try:
            assert isinstance(animate, bool)
        except AssertionError:
            raise TypeError("animate must be a boolean value")
        if self.visual_type == "scientific":
            self.__scientific_plot(animate=animate)
        else:
            self.__animation()    

Methods

def visualize(self, animate: bool = False) ‑> None

Method of the visualization class that generates the visualization for the trajectories based on the visualization type

Args

animate : bool
Used for scientific visualization and control whether an animation is generated. Defaults to False
dense_body_sim : bool
Used for scientific visualization and control whether the visualization is for a dense body simulation. Defaults to False

Raises

TypeError
When the animate not a boolean value.