A very crude galaxy simulator.

newtonian-gravity.py

"""

Code by:
Brian Schrader
09-21-2024
"""

from datetime import datetime
import logging
import time

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from matplotlib import colors


logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


G = -6.6743 #* 10e-11


np.seterr(divide='ignore')


def A(P_n, G, M):
    dP = P_n - P_n[:,None]
    R = np.linalg.norm(dP, axis=-1)

    a_R = G * M / (R ** 2)
    a_R[np.isnan(a_R)] = 0

    a_X = (dP.T / (R * a_R)).T
    a_X[np.isnan(a_X)] = 0
    return np.sum(a_X, axis=0)


def P_V(P_n1, V_n1, G, M, dt=0.5):
    a_n = A(P_n1, G, M)
    V_n = V_n1 + a_n * dt
    P_n = P_n1 + V_n * dt
    return P_n, V_n


def plot(Ps, limit, anim_step=1):
    logger.info('Plotting...')
    background_color = 0.1, 0.1, 0.13, 1
    dots_color = 1, 1, 1, 1

    fig, ax = plt.subplots(subplot_kw={"projection": "3d"}, facecolor=background_color)
    ax.set_facecolor(background_color)
    ax.tick_params(labelcolor=background_color)

    ax.set_xlim([-limit, limit])
    ax.set_ylim([-limit, limit])
    ax.set_zlim([-limit, limit])
    ax.grid(False)

    ax.get_xaxis().set_pane_color(background_color)
    ax.get_yaxis().set_pane_color(background_color)
    ax.get_zaxis().set_pane_color(background_color)

    def update(args):
        t, P = args
        ax.clear()
        ax.grid(False)

        ax.get_xaxis().set_pane_color(background_color)
        ax.get_yaxis().set_pane_color(background_color)
        ax.get_zaxis().set_pane_color(background_color)

        ax.view_init(elev=(t/5)%90000, azim=(t/5)%360000)

        ax.set_xlim([-limit, limit])
        ax.set_ylim([-limit, limit])
        ax.set_zlim([-limit, limit])

        ax.set_title(f'{t=}', color='white')
        plot = ax.scatter(P[:,0], P[:,1], P[:,2], s=1, color=dots_color)
        return plot

    ani = animation.FuncAnimation(
        fig=fig,
        func=update,
        frames=list(enumerate(Ps))[::anim_step],
        interval=100,
    )
    ts = datetime.now().strftime('%Y-%M-%d-%H-%M')
    ani.save(f'simulation-{ts}.m4v')
    plt.close()


def simulate(
    P_0,
    V_0,
    M,
    steps,
    n_nodes,
    args=None,
    update_interval=10,
):
    logger.info('Running...')
    Ps = [P_0]
    P_prev = P_0
    V_prev = V_0
    t = 0
    start = t0 = time.time()
    for t in range(steps):
        P_t, V_t = P_V(P_prev, V_prev, G, M)
        P_prev = P_t
        V_prev = V_t
        if t % update_interval == 0:
            logger.info(f'{t=}: d[{update_interval}] = {(time.time()-t0):.2f}s')
            Ps.append(P_t)
            t0 = time.time()

    logger.info(
        f'Simulation complete! ({t=} steps) - total time = {(time.time()-start):.2f}s'
    )
    return Ps


def main(n_nodes=1_200, p_scale=1e8, steps=1_000_000, v_scale=1e2, m_scale=2e12):
    logger.info(f'Settings: {n_nodes=}, {p_scale=}, {v_scale=}, {m_scale=}, {steps=}')
    P_0 = np.random.randint(-p_scale, p_scale, size=(n_nodes, 3))
    V_0 = np.random.randint(-v_scale, v_scale, size=(n_nodes, 3))
    M = m_scale * np.random.uniform(0.7, 10, size=(n_nodes, 1))
    Ps = simulate(P_0, V_0, M, steps, n_nodes)
    plot(Ps, p_scale * 2)


if __name__ == '__main__':
    main()