mirror of
https://github.com/Findus23/rebound-collisions.git
synced 2024-09-19 15:53:48 +02:00
66 lines
2.5 KiB
Python
66 lines
2.5 KiB
Python
|
from os.path import expanduser
|
||
|
|
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
import numpy as np
|
||
|
|
from matplotlib.axes import Axes
|
||
|
|
from matplotlib.cm import ScalarMappable
|
||
|
|
from matplotlib.colors import Normalize
|
||
|
|
from matplotlib.figure import Figure
|
||
|
|
from rebound import SimulationArchive
|
||
|
|
from scipy.constants import mega
|
||
|
|
|
||
|
|
from extradata import ExtraData
|
||
|
|
from utils import filename_from_argv, get_water_cmap
|
||
|
|
|
||
|
|
mean_mass = 5.208403167890638e+24 # constant between plots
|
||
|
|
size_mult = 100
|
||
|
|
|
||
|
|
|
||
|
|
def plot_file(file, time, ax: Axes, mode):
|
||
|
|
fn = filename_from_argv(file)
|
||
|
|
sa = SimulationArchive(str(fn.with_suffix(".bin")))
|
||
|
|
ed = ExtraData.load(fn)
|
||
|
|
sim = sa.getSimulation(t=time)
|
||
|
|
|
||
|
|
water_fractions = [ed.pd(p).water_mass_fraction for p in sim.particles[1:]]
|
||
|
|
a = [p.a for p in sim.particles[1:]]
|
||
|
|
e = [p.e for p in sim.particles[1:]]
|
||
|
|
m = np.array([p.m for p in sim.particles[1:]])
|
||
|
|
# m[:2] /= 1e2 # reduce size of gas giants
|
||
|
|
sizes = (np.array(m) / mean_mass) ** (2 / 3) * size_mult
|
||
|
|
with np.errstate(divide='ignore'): # allow 0 water (becomes -inf)
|
||
|
|
color_val = (np.log10(water_fractions) + 5) / 5
|
||
|
|
cmap = get_water_cmap()
|
||
|
|
colors = cmap(color_val)
|
||
|
|
ax.set_xlim(0, 5.5)
|
||
|
|
ax.set_ylim(-0.05, .8)
|
||
|
|
ax.text(0.05, 0.93, f"{time / mega:.0f} Myr", size=10,
|
||
|
|
horizontalalignment='left', verticalalignment='top', transform=ax.transAxes)
|
||
|
|
ax.text(0.95, 0.93, mode, size=10,
|
||
|
|
horizontalalignment='right', verticalalignment='top', transform=ax.transAxes)
|
||
|
|
ax.scatter(a, e, s=sizes, zorder=10, c=colors)
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
fig: Figure
|
||
|
|
modes = ["RBF", "NN", "PM", "LZ"]
|
||
|
|
fig, axes = plt.subplots(5, 4, sharex=True, sharey=True, figsize=(10, 13))
|
||
|
|
for col_nr, file in enumerate(
|
||
|
|
["data/final_rbf_1.bin", "data/final_nn_1.bin", "data/final_pm_1.bin", "data/final_lz_1.bin"]):
|
||
|
|
for row_nr, time in enumerate([1 * mega, 5 * mega, 20 * mega, 50 * mega, 100 * mega]):
|
||
|
|
plot_file(file, time, axes[row_nr][col_nr], mode=modes[col_nr])
|
||
|
|
fig_colorbar = plt.figure()
|
||
|
|
fig_colorbar.colorbar(ScalarMappable(norm=Normalize(vmin=-5, vmax=0), cmap=get_water_cmap()), aspect=20,
|
||
|
|
label="log(water mass fraction)", orientation="horizontal")
|
||
|
|
fig.supxlabel("semi-major axis [AU]")
|
||
|
|
fig.supylabel("eccentricity")
|
||
|
|
fig.tight_layout()
|
||
|
|
fig_colorbar.gca().remove()
|
||
|
|
fig.savefig(expanduser(f"~/tmp/timeplot.pdf"), dpi=300)
|
||
|
|
fig_colorbar.savefig(expanduser(f"~/tmp/timeplot_colorbar.pdf"), bbox_inches='tight')
|
||
|
|
plt.show()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main()
|