odisseo.dynamics module

odisseo.dynamics.direct_acc(state: Array, mass: Array, config: SimulationConfig, params: SimulationParams, return_potential=False)

Compute acceleration of all particles due to all other particles by vmap of the single_body_acc function.

Parameters:

• state – Array of shape (N, 2, 3) containing the positions and velocities of the particles.

• mass – Array of shape (N,) containing the masses of the particles.

• config – Configuration object containing the number of particles (N_particles) and softening parameter.

• params – Parameters object containing the gravitational constant (G).

• return_potential – If True, also return the potential energy. Defaults to False.

Returns:

Array of shape (N, 3) containing the accelerations of the particles. Array of shape (N,) containing the potential energy of the particles, if return_potential is True.

odisseo.dynamics.direct_acc_for_loop(state: Array, mass: Array, config: SimulationConfig, params: SimulationParams, return_potential: bool = False) → Array | tuple[Array, Array]

Compute the direct acceleration matrix for a system of particles. Uses a double for loop and Newton’s third low to reduce the computation from O(N^2) to O(N^2 /2).

Parameters:

• state – Array of shape (N, 2, 3) containing the positions and velocities of the particles.

• mass – Array of shape (N,) containing the masses of the particles.

• config – Configuration object containing the number of particles (N_particles) and softening parameter.

• params – Parameters object containing the gravitational constant (G).

• return_potential – If True, also return the potential energy. Defaults to False.

Returns:

Array of shape (N, 3) containing the accelerations of the particles. Array of shape (N,) containing the potential energy of the particles, if return_potential is True.

odisseo.dynamics.direct_acc_laxmap(state: Array, mass: Array, config: SimulationConfig, params: SimulationParams, return_potential=False)

Compute acceleration of all particles due to all other particles by using lax.map of the single_body_acc function. If config.double_map is True, lax.map uses lax.map for both loops, otherwise the inner loop is vectorized using vmap. Memory usage is reduced by using lax.map instead of vmap thanks to batching.

Parameters:

• state – Array of shape (N, 2, 3) containing the positions and velocities of the particles.

• mass – Array of shape (N,) containing the masses of the particles.

• config – Configuration object containing the number of particles (N_particles) and softening parameter.

• params – Parameters object containing the gravitational constant (G).

• return_potential – If True, also return the potential energy. Defaults to False.

Returns:

Array of shape (N, 3) containing the accelerations of the particles. Array of shape (N,) containing the potential energy of the particles, if return_potential is True.

odisseo.dynamics.direct_acc_matrix(state: Array, mass: Array, config: SimulationConfig, params: SimulationParams, return_potential: bool = False) → Array | tuple[Array, Array]

Compute the direct acceleration matrix for a system of particles. Uses matrix operations.

Parameters:

• state – Array of shape (N, 2, 3) containing the positions and velocities of the particles.

• mass – Array of shape (N,) containing the masses of the particles.

• config – Configuration object containing the number of particles (N_particles) and softening parameter.

• params – Parameters object containing the gravitational constant (G).

• return_potential – If True, also return the potential energy. Defaults to False.

Returns:

Array of shape (N, 3) containing the accelerations of the particles. Array of shape (N,) containing the potential energy of the particles, if return_potential is True.

odisseo.dynamics.direct_acc_sharding(state: Array, mass: Array, config: SimulationConfig, params: SimulationParams, return_potential: bool = False) → Array | tuple[Array, Array]

Compute the direct acceleration matrix for a system of particles. Shard the positions to allow for parallel computation. CURRENTLY NOT WORKING.

Parameters:

• state – Array of shape (N, 2, 3) containing the positions and velocities of the particles.

• mass – Array of shape (N,) containing the masses of the particles.

• config – Configuration object containing the number of particles (N_particles) and softening parameter.

• params – Parameters object containing the gravitational constant (G).

• return_potential – If True, also return the potential energy. Defaults to False.

Returns:

Array of shape (N, 3) containing the accelerations of the particles. Array of shape (N,) containing the potential energy of the particles, if return_potential is True.

odisseo.dynamics.no_self_gravity(state: Array, mass: Array, config: SimulationConfig, params: SimulationParams, return_potential=False)

Remove the self interaction between particles.

Parameters:

• state – Array of shape (N, 2, 3) containing the positions and velocities of the particles.

• mass – Array of shape (N,) containing the masses of the particles.

• config – Configuration object containing the number of particles (N_particles) and softening parameter.

• params – Parameters object containing the gravitational constant (G).

• return_potential – If True, also return the potential energy. Defaults to False.

Returns:

Array of shape (N, 3) containing the accelerations of the particles. Array of shape (N,) containing the potential energy of the particles, if return_potential is True.

odisseo.dynamics.single_body_acc(particle_i: Array, particle_j: Array, mass_i: Array, mass_j: Array, config: SimulationConfig, params: SimulationParams) → tuple[Array, Array]

Compute acceleration of particle_i due to particle_j.

Parameters:

• particle_i – Position and velocity of particle_i.

• particle_j – Position and velocity of particle_j.

• mass_i – Mass of particle_i.

• mass_j – Mass of particle_j.

• config – Configuration parameters.

• params – Simulation parameters.

Returns:

The acceleration of particle_i due to particle_j, and the potential felt particle_i due to particle_j.