21 add celestial nbody example

[JaLuka98]

This pull request adds a new example script, examples/nbody_gravity.py, which demonstrates an N-body gravitational simulation using JAX and matplotlib. The script includes functionality for simulating the motion of celestial bodies, analyzing energy conservation, and visualizing the results as both static plots and an animated GIF.

Key changes:

New N-body simulation example:

• Added examples/nbody_gravity.py, which simulates the gravitational interactions of four celestial bodies (Sun, Earth, Mars, and a comet) using a leapfrog integrator. The script leverages JAX for numerical computation and matplotlib for visualization.

Energy diagnostics and visualization:

• Implemented energy diagnostics to compute and plot the total energy (kinetic + potential) over time, verifying energy conservation in the simulation.

Animation of the simulation:

• Created an animated visualization of the N-body trajectories, saved as a GIF (nbody_simulation.gif), using matplotlib's FuncAnimation and PillowWriter.