odisseo.integrators module

odisseo.integrators module#

odisseo.integrators.RungeKutta4(state: Array, mass: Array, dt: Shaped[Array, ''], config: SimulationConfig, params: SimulationParams)[source]#

Simple implementation of a 4th order Runge-Kutta integrator for N-body simulations.

Parameters:

  • state (jax.numpy.ndarray) – The state of the particles, where the first column represents positions and the second column represents velocities.

  • mass (jax.numpy.ndarray) – The mass of the particles.

  • dt (float) – Time-step for current integration.

  • config (object) – Configuration object containing the acceleration scheme and external accelerations.

  • params (dict) – Additional parameters for the acceleration functions.

Returns:

jax.numpy.ndarray – The updated state of the particles.

odisseo.integrators.diffrax_solver(state: Array, mass: Array, dt: Shaped[Array, ''], config: SimulationConfig, params: SimulationParams) → Array[source]#

Diffrax backhend

Parameters:

  • state (jax.numpy.ndarray) – The state of the particles, where the first column represents positions and the second column represents velocities.

  • mass (jax.numpy.ndarray) – The mass of the particles.

  • dt (float) – Time-step for current integration.

  • config (object) – Configuration object containing the acceleration scheme and external accelerations.

  • params (dict) – Additional parameters for the acceleration functions.

Returns:

jax.numpy.ndarray – The updated state of the particles.

odisseo.integrators.leapfrog(state: Array, mass: Array, dt: Shaped[Array, ''], config: SimulationConfig, params: SimulationParams)[source]#

Simple implementation of a symplectic Leapfrog (Verlet) integrator for N-body simulations.

Parameters:

  • state (jax.numpy.ndarray) – The state of the particles, where the first column represents positions and the second column represents velocities.

  • mass (jax.numpy.ndarray) – The mass of the particles.

  • dt (float) – Time-step for current integration.

  • config (object) – Configuration object containing the acceleration scheme and external accelerations.

  • params (dict) – Additional parameters for the acceleration functions.

Returns:

jax.numpy.ndarray – The updated state of the particles.