Asteroid derivatives snippet

asteroid.py

def calculate_derivs(positions, t, m_sun, m_planet, G):
    # Unpack positions and velocities
    asteroid_x = positions[0]
    asteroid_vx = positions[1]
    asteroid_y = positions[2]
    asteroid_vy = positions[3]

    sun_x = positions[0]
    sun_vx = positions[1]
    sun_y = positions[2]
    sun_vy = positions[3]

    planet_x = positions[0]
    planet_vx = positions[1]
    planet_y = positions[2]
    planet_vy = positions[3]

    # Calculate Distances
    r_sun_asteroid = np.sqrt((asteroid_x - sun_x)**2 + (asteroid_y - sun_y)**2)
    r_sun_planet = ...
    # Calculate accelerations
    # For asteroid, take into account sun and planet
    asteroid_ax = ...
    asteroid_ay = ...
    # For sun, only take into account planet
    sun_ax = G * m_planet * (1 / r_sun_planet**3) * (planet_x - sun_x)
    sun_ay = ...
    # For planet, only take into account sun
    planet_ax = ...
    planet_ay = ...
    
    # Return the time derivatives of the input positions and velocities
    return [asteroid_vx, asteroid_ax, asteroid_vy, asteroid_ay...