def plane_angle(sol, axis):
    # Calculate angle of separation of first two particles in (sol) relative to (axis)

    D = 3
    unit = zeros(D)
    unit[axis] = 1

    x1s = stack_sol(sol, D)[:, :D, 0]
    x2s = stack_sol(sol, D)[:, :D, 1]
    rs = x2s - x1s
    Rs = mag(rs, axis=1)

    dotted = np.dot(rs, unit)

    return np.arccos(dotted / Rs)
def circle_errors(sol, m, R, t):
    # Differences between computed positions of (sol) and analytic circular orbit (R) about (m) at (t)

    D = 2  # sol assumed 2D
    xs = stack_sol(
        sol,
        D)[:, :D,
           1].transpose()  # extract position of second particle at for each t
    difs = xs - analytic_circle(m, R, t)
    return mag(difs, axis=0)
def animate_3D(sol, n_frames, ax_lims, filename):
    """
     Save 3D animation of (sols) [assumed 3D] as (filename).html
     within (ax_lims) over (n_frames)
    """
    # extract x, y, z coordinates of system over whole simulation
    D = 3
    p = stack_sol(sol, D)
    xs, ys, zs = p[:, 0], p[:, 1], p[:,2]
    
    # get axes limits (or create them if not specified)
    if np.size(ax_lims) == 0:
        # capture all points with least whitespace possible
        xmax = ymax = zmax = np.max(np.abs([xs, ys, zs]))
        xmin = ymin = zmin = - xmax
    else:
        xmin, xmax, ymin, ymax, zmin, zmax = ax_lims
        
    
    plt.close("all")
    fig = plt.figure()
    
    # set axes
    ax = fig.add_subplot(111, projection='3d')
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(ymin, ymax)
    ax.set_zlim(zmin, zmax)
    ax.set_xlabel('x/m', fontsize=12)
    ax.set_ylabel('y/m', fontsize=12)
    ax.set_zlabel('z/m', fontsize=12)
    plt.gca().set_aspect('equal', adjustable='box')    # ensure equal scaling of axes
    fr, = ax.plot([], [], 'bo', markersize=0.1)
    
    # plot path over time of first two particles
    ax.plot(xs[:,0], ys[:,0], zs[:,0], 'r-', linewidth = 0.3)
    ax.plot(xs[:,1], ys[:,1], zs[:,1], 'r-', linewidth = 0.3)
    
    inter = int( sol.shape[0] / n_frames )    # number of steps between each frame
    
    # function specifying how a given frame is generated
    def next_frame(i):
        # update coords of each particle
        j = inter*i
        fr.set_data(xs[j], ys[j])
        fr.set_3d_properties(zs[j])
        return fr
    
    # animate and save
    ani = animation.FuncAnimation(fig, next_frame, frames=n_frames, interval=20)

    ani.save(filename + '.html')
    plt.close()
def snapshot_parabola(sol, step, t, A, num_paths, ax_lims, n_D, filename):
    """
     Plot snapshot of particles in (sol) at frame (step) in time (t)
     (A) is radius of first two particles' halos (set to 0 if point mass)
     Plot within (ax_lims) and save as (filename).pdf
     - Also plots path of first (num_paths) particles over t
     - projects particles into xy-plane if n_D > 2
    """
    D = 2
    
    p = stack_sol(sol, n_D)    # xs and ys over all time
    xs, ys = p[:, 0], p[:, 1]
    
    X, Y = xs[step, :2], ys[step, :2]    # centres
    test_x, test_y = xs[step, 2:], ys[step, 2:]    # test particles
    
    if np.size(ax_lims) == 0:
        rb = 1.25 # plot captures all existing points with 100(rb-1)% whitespace on all sides
        min1 = min2 =  rb * np.min([xs[step], ys[step]])
        max1 = max2 = rb * np.max([xs[step], ys[step]])
    else:
        min1, max1, min2, max2 = ax_lims
    
    fig, ax = plt.subplots()
    
    plt.gca().set_aspect('equal', adjustable='box')    # ensure equal scaling of axes
    plt.xlim(min1, max1)
    plt.ylim(min2, max2)
    ax.scatter(X, Y, c='r', s=20)
    ax.scatter(test_x, test_y, c='b', s=0.1)
    
    # plot path of first num_paths particles
    for i in range(num_paths):
        ax.plot(xs[:,i], ys[:,i], 'r', linewidth=0.5)
    
    # plot halo of radius A around first two particles
    circle1 = plt.Circle((X[0], Y[0]), 0.2, facecolor='none', edgecolor='r', linewidth=0.9, linestyle='--')
    circle1.set_radius(A)
    ax.add_artist(circle1)
    
    circle2 = plt.Circle((X[1], Y[1]), 0.2, facecolor='none', edgecolor='r', linewidth=0.9, linestyle='--')
    circle2.set_radius(A)
    ax.add_artist(circle2)
    
    plt.title('t = ' + str(round(t[step], 1)) + 's')
    plt.xlabel('x/m')
    plt.ylabel('y/m')
    plt.savefig(filename + '.pdf')
    plt.show()
def snapshot_3D(sol, step, t, ax_lims, filename):
    """
     Plot 3D snapshot of particles in (sol) at frame (step) in time (t)
     Plot within (ax_lims) and save as (filename).pdf
     - sol is assumed to be 3D
    """
    D = 3
    
    p = stack_sol(sol, D)    # xs, ys, zs over all time
    xs, ys, zs = p[:, 0], p[:, 1], p[:, 2]
    
    X, Y, Z = xs[step, :2], ys[step, :2], zs[step, :2]    # centres
    test_x, test_y, test_z = xs[step, 2:], ys[step, 2:], zs[step, 2:]    # test particles
    
    # get ax_lims or generate if none given
    if np.size(ax_lims) == 0:
        # capture all points with least whitespace possible
        xmax = ymax = zmax = np.max(np.abs([xs, ys, zs]))
        xmin = ymin = zmin = - xmax
    else:
        xmin, xmax, ymin, ymax, zmin, zmax = ax_lims
    
    plt.close("all")
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')    # 3D plot
    
    # set axes
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(ymin, ymax)
    ax.set_zlim(zmin, zmax)
    ax.set_xlabel('x/m', fontsize=12)
    ax.set_ylabel('y/m', fontsize=12)
    ax.set_zlabel('z   ', fontsize=12)
    plt.tick_params(labelsize=10)
    
    ax.scatter(X, Y, Z, c='r', s=10)
    ax.scatter(test_x, test_y, test_z, c='b', s=0.8)
    
    plt.gca().set_aspect('equal', adjustable='box')    # ensure equal scaling of axes
    
    # plot path of first two particles
    ax.plot(xs[:,0], ys[:,0], zs[:,0], 'r', linewidth=0.4)
    ax.plot(xs[:,1], ys[:,1], zs[:,1], 'r', linewidth=0.4)
    
    plt.title('t = ' + str(round(t[step], 1)) + 's')
    plt.savefig(filename + '.pdf')
    plt.show()
def animate_2D(sol, n_frames, ax_lims, n_D, filename):
    """
     Save animation of (n_D)-dimensional (sols) as (filename).html
     within (ax_lims) over (n_frames)
     - if n_D > 2, the system is projected into the xy-plane
    """
    # extract x and y coordinates of system over whole simulation
    p = stack_sol(sol, n_D)
    xs, ys = p[:, 0], p[:, 1]
    
    plt.close("all")
    fig, ax = plt.subplots()
    
    # get axes limits (or create them if not specified)
    if np.size(ax_lims) == 0:
        # capture all existing points with at least 100(rb-1)% whitespace on all sides
        rb = 1.25
        min1 = min2 =  rb * np.min([xs, ys])
        max1 = max2 = rb * np.max([xs, ys])
    else:
        min1, max1, min2, max2 = ax_lims
    
    # set axes
    plt.xlim(min1, max1)
    plt.ylim(min2, max2)
    plt.xlabel('x/m')
    plt.ylabel('y/m')
    plt.gca().set_aspect('equal', adjustable='box')    # ensure equal scaling of axes
    fr, = ax.plot([], [], 'bo', markersize=0.1)
    
    inter = int( sol.shape[0] / n_frames )    # number of steps between each frame
    
    # function specifying how a given frame is generated
    def next_frame(i):
        # update coords of each particle
        j = inter*i
        fr.set_data(xs[j], ys[j])
        return fr
    
    # animate and save
    ani = animation.FuncAnimation(fig, next_frame, frames=n_frames, interval=20)

    ani.save(filename + '.html')
    plt.close()
def snapshot_projection(sol, axis, step, t, ax_lims, filename):
    """
     Plot 2D snapshot of 3D (sol) after removing components parallel to (axis)
     - e.g. axis=1 gives projection in yz plane
    """
    D = 3
    
    p = stack_sol(sol, D)    
    pro_p = np.delete(p, axis,  1)    # remove axis component of position
    x1s, x2s = pro_p[:, 0], pro_p[:, 1]    # x1s and x2s over all time
    
    X1, X2 = x1s[step, :2], x2s[step, :2]    # centres
    test_x1, test_x2 = x1s[step, 2:], x2s[step, 2:]    # test particles
    
    if np.size(ax_lims) == 0:
        rb = 1.25 # plot captures all existing points with 100(rb-1)% whitespace on all sides
        min1 = min2 =  rb * np.min([x1s[step], x2s[step]])
        max1 = max2 = rb * np.max([x1s[step], x2s[step]])
    else:
        min1, max1, min2, max2 = ax_lims
    
    fig, ax = plt.subplots()
    
    plt.gca().set_aspect('equal', adjustable='box')    # ensure equal scaling of axes
    plt.xlim(min1, max1)
    plt.ylim(min2, max2)
    ax.scatter(X1, X2, c='r', s=20)
    ax.scatter(test_x1, test_x2, c='b', s=0.1)
    
    plt.title('t = ' + str(round(t[step], 1)) + 's')
    
    # plot axes titles depending on projection plane
    if axis==0:
        plt.xlabel('y/m')
        plt.ylabel('z/m')
    elif axis==1:
        plt.xlabel('x/m')
        plt.ylabel('z/m')
    else:
        plt.xlabel('x/m')
        plt.ylabel('y/m')
    
    plt.savefig(filename + '.pdf')
    plt.show()
def snapshot_circle(sol, step, t, A, ax_lims, n_D, filename):
    """
     Plot snapshot of particles in (sol) at frame (step) in time (t)
     (A) is radius of first particle halo (set to 0 if point mass)
     Plot within (ax_lims) and save as (filename).pdf
     - Also plots path of all particles apart from the first over t
     - projects particles into xy-plane if n_D > 2
    """
    D = 2
    
    p = stack_sol(sol, n_D)
    xs, ys = p[:, 0], p[:, 1]    # xs and ys over all time
    
    # xs and ys at step
    X, Y = xs[step, 0], ys[step, 0]    # centre
    test_x, test_y = xs[step, 1:], ys[step, 1:]    # test particles
    
    # get axes limits (or create them if not specified)
    if np.size(ax_lims) == 0:
        rb = 1.25 # plot captures all existing points with 100(rb-1)% whitespace on all sides
        min1 = min2 =  rb * np.min([xs, ys])
        max1 = max2 = rb * np.max([xs, ys])
    else:
        min1, max1, min2, max2 = ax_lims
    
    fig, ax = plt.subplots()
    
    plt.gca().set_aspect('equal', adjustable='box')    # ensure equal scaling of axes
    plt.xlim(min1, max1)
    plt.ylim(min2, max2)
    ax.scatter(X, Y, c='r', s=40)
    ax.scatter(test_x, test_y, c='b', s=20)
    ax.plot(xs[:,1:], ys[:,1:], 'b', linewidth=0.1)    # plot path of all but first particle
    
    # plot halo of radius A around first particle
    circle1 = plt.Circle((X, Y), 0.2, facecolor='none', edgecolor='r', linewidth=0.5)
    circle1.set_radius(A)
    ax.add_artist(circle1)
    
    plt.title('t = ' + str(round(t[step], 1)) + 's')
    plt.xlabel('x/m')
    plt.ylabel('y/m')
    plt.savefig(filename + '.pdf')
    plt.show()
plt.xlabel('t/s')
plt.ylabel('Relative error')
plt.savefig('Relative error for circle R=2.pdf')
plt.show()

plt.plot(Rs, max_errs)
plt.xlabel('Radius of orbit/m')
plt.ylabel('Maximum error/m')
plt.savefig('Max error of circular orbit vs R.pdf')
plt.show()

# compare simulation and theory for R=2 orbit
D = 2
i = 0  # index of R=2 orbit
R = Rs[i]
xs = stack_sol(circle_sols[i], D)[:, 0,
                                  1].transpose()  # extract xs of test particle
analytic_xs = analytic_circle(M, R, t)[0]

start = int((99 / 100) * len(t))  # begin plot at start of 100th orbit
plt.plot(t[start:], xs[start:], label='computed')
plt.plot(t[start:], analytic_xs[start:], 'k:', label='analytical')
plt.legend(loc='best')
plt.xlabel('t/s')
plt.ylabel('x component of position/m')
plt.savefig('comparison of computed and analytic circular motion.pdf')
plt.show()

step = len(t) - 1
A = 0
snapshot_circle(circle_sols[i], step, t, A, [], D, 'snapshot - circle R=2')
galaxy = np.append(halo, stars, axis=1)
q0 = galaxy.flatten()

g_halo_A = lambda m, x, x1s: g_halo(m, A, x, x1s)

halo_circle_sol = odeint(g_diffs, q0, t, args=([M], g_halo_A, D))

snap_start = int((end_time - ((3 / 4) * T_halo)) /
                 step_size)  # step no. for 1/4 way through last orbit
snap_end = len(halo_circle_sol[snap_start:]) - 1
snapshot_circle(halo_circle_sol[snap_start:], snap_end, t[snap_start:], A, [],
                D, 'snapshot - halo test')

# compare simulation and theory
xs = stack_sol(halo_circle_sol,
               D)[:, 0, 1].transpose()  # extract xs of test particle
analytic_xs = analytic_halo(M, R, t)[0]

start = int((99 / 100) * len(t))  # begin plot at start of 100th orbit
plt.plot(t[start:], xs[start:], label='computed')
plt.plot(t[start:], analytic_xs[start:], 'k:', label='analytical')
plt.legend(loc='best')
plt.xlabel('t/s')
plt.ylabel('x component of position/m')
plt.savefig('comparison of computed and analytic circular motion.pdf')
plt.show()

# plot errors for R=2
h_errs = halo_errors(halo_circle_sol, M, Rs[0], t)

plt.plot(t, h_errs)
# snapshots
no_rotation_sol = np.loadtxt('3D collision solutions - standard.txt')
step = len(no_rotation_sol) - 1
inner_lims_3D = [-20, 20, -20, 20, -20, 20]
t_snap = t[0::interval]
snapshot_3D(no_rotation_sol, step, t_snap, inner_lims_3D,
            '3D snapshot - no rotation')

step = 149
for i, n in enumerate(ns):
    snapshot_3D(rotated_x_sols[i], step, t_snap, inner_lims_3D,
                '3D snapshot - x rotation, n = ' + str(n))
    snapshot_3D(rotated_y_sols[i], step, t_snap, inner_lims_3D,
                '3D snapshot - y rotation, n = ' + str(n))

# 2D projections of n=1 x rotation
step = 149
A = 0
helix_sol = rotated_x_sols[0]

# to find lims which place galaxy centre at origin for each of 3 projections
centre = stack_sol(helix_sol, D)[step, :D, 0]
individual_lims = [[x - 20, x + 20] for x in centre]
project_lims = [
    np.concatenate(np.delete(individual_lims, i, 0)) for i in range(D)
]

axes = np.arange(3)
for axis in axes:
    snapshot_projection(helix_sol, axis, step, t_snap, project_lims[axis],
                        'snapshot projection - helix ' + str(axis))