Solve equation of movement including shell crossing

nfw.py

import numpy as np
from astropy import constants as C
from tqdm import tqdm
import os
import h5py


N = 10000
nsteps = 1_000
iout = 10


def do_dump(i):
    return i % iout == 0


def dump(i, R, M, v, order):
    fname = os.path.join('output', f'{i:05d}.h5')
    with h5py.File(fname, 'w') as f:
        f['R'] = R
        f['M'] = M
        f['v'] = v
        f['order'] = order


def main():
    R0 = np.geomspace(1e-4, 1, N)
    M0 = 4*np.pi*R0**2*np.gradient(R0)

    R = R0.copy()
    M = M0.copy()
    v = np.zeros(N) # (np.random.rand(N) - 0.5) / 100
    rho = np.zeros(N)

    # Temp quantities
    Mo = np.zeros(N)
    vo = np.zeros(N)
    Ro = np.zeros(N)
    Mino = np.zeros(N)
    Fo = np.zeros(N)

    # Gravitational constant
    G = C.G.cgs.value
    history = []
    for i in tqdm(range(nsteps)):
        order = np.argsort(R)
        Mo[:] = M[order]
        vo[:] = v[order]
        Ro[:] = R[order]

        # Compute inside mass
        Mino[:] = np.cumsum(Mo) - Mo

        # Compute acceleration of shell [force / m]
        Fo[:] = -G * Mino / Ro**2

        # Compute timestep [|R / a|**(1/2)]
        dt = np.sqrt(np.abs(Ro / Fo)).min() / 100

        # Evolve velocity
        v[order] += dt*Fo

        # Evolve radii
        R[order] += dt*v

        # Spherical symmetry R = |R|
        R = np.abs(R)

        # Compute rho
        rho[order] = Mino / (4*np.pi/3*Ro**3)

        if do_dump(i):
            dump(i, R, M, v, order)


if __name__ == '__main__':
    os.makedirs('output', exist_ok=True)
    main()