Commit d59e2678 authored by Zoheyr Doctor's avatar Zoheyr Doctor 🔭
Browse files

made script to create individual gifs for each event and generated the gifs

parent abd98ea0
#!/usr/bin/env python
import numpy as np
from twirl.binary import Binary
from matplotlib import pyplot as plt
from matplotlib import animation
from matplotlib.patches import Circle
def make_parser():
from argparse import ArgumentParser
parser = ArgumentParser()
parser.add_argument(
"file",
type=str,
help='json file with median values',
)
parser.add_argument(
"outfolder",
type=str,
default='.',
help='folder to output gifs'
)
parser.add_argument(
"--round",
help="round orbital frequencies so that the movie can seamlessly loop",
action="store_true"
)
parser.add_argument(
"--dark",
help="use dark background",
action="store_true"
)
return parser
def main():
import json
import sys
import sys
# parse arguments
parser = make_parser()
args = parser.parse_args()
with open(args.file,'r') as json_file:
data = json.load(json_file)
# get number of events
events = list(data.keys())
n_events = len(events)
# make axes
if args.dark:
plt.style.use('dark_background')
# set frame rate and movie length
fps = 30. # frames per second
length = 10. # seconds
# determine orbital frequencies
if args.round:
omegas = np.array([1./(data[event]['mass_1_source']['med']
+ data[event]['mass_2_source']['med']
) for event in events
])
omegas = np.round(omegas/np.amin(omegas),0)
omegas = {event : omegas[i]/length for i,event in enumerate(events)}
else:
omegas = {event : 20./(data[event]['mass_1_source']['med']
+ data[event]['mass_2_source']['med']
) for event in events
}
pos_dict = {} # dictionary for all compact object positions
binary_dict = {} # dictionary for all Binary objects
# loop over all events
for event in events:
print('evolving ',event)
event_pos_dict = {} # dictionary for positions of this event
m1 = data[event]['mass_1_source']['med']
m2 = data[event]['mass_2_source']['med']
# create a Binary instance for the event
binary = Binary(
m1 = m1,
m2 = m2,
c1 = 'k',
c2 = 'k',
alpha = 45.,
beta = 20.,
gamma = 80.,
omega = omegas[event],
frame_rate = fps
)
# evolve the binary over one orbit
n_frames = int(length*fps)
# populate positions for event
event_pos_dict = {
var: np.empty(n_frames)
for var in ['x1','y1','z1','x2','y2','z2']
}
for i in range(n_frames):
binary.orbit()
pos1,pos2 = binary.project()
event_pos_dict['z1'][i] = pos1[0]
event_pos_dict['z2'][i] = pos2[0]
event_pos_dict['x1'][i] = pos1[1]
event_pos_dict['y1'][i] = pos1[2]
event_pos_dict['x2'][i] = pos2[1]
event_pos_dict['y2'][i] = pos2[2]
# write out evolution to dictionary
pos_dict[event] = event_pos_dict
binary_dict[event] = binary
for ei,event in enumerate(events):
fig,ax = plt.subplots()
ax.set_axis_off()
artist_dict = {} # dictionary for matplotlib artists
def init():
# create artists for both stars
artist_dict['%s_m1'%event] = ax.add_artist(
Circle((pos_dict[event]['x1'][0],pos_dict[event]['y1'][0]),
radius=binary_dict[event].m1/200,
edgecolor='white',
facecolor='black'
)
)
artist_dict['%s_m2'%event] = ax.add_artist(
Circle((pos_dict[event]['x2'][0],pos_dict[event]['y2'][0]),
radius=binary_dict[event].m2/200,
edgecolor='white',
facecolor='black'
)
)
# make sure star that is closer is shown on top of the other
zorder_1 = int(pos_dict[event]['z1'][0]>pos_dict[event]['z2'][0])
zorder_2 = ~zorder_1
artist_dict['%s_m1'%event].set_zorder(zorder_1)
artist_dict['%s_m2'%event].set_zorder(zorder_2)
ax.set_xlim([-1,1])
ax.set_ylim([-1,1])
ax.set_title(event)
return artist_dict
def animate(i):
# update positions of stars
artist_dict['%s_m1'%event].set_center((pos_dict[event]['x1'][i],pos_dict[event]['y1'][i]))
artist_dict['%s_m2'%event].set_center((pos_dict[event]['x2'][i],pos_dict[event]['y2'][i]))
# make sure star that is closer is shown on top of the other
if pos_dict[event]['z1'][i]>pos_dict[event]['z2'][i]:
zorder_1 = 2
zorder_2 = 1
else:
zorder_1 = 1
zorder_2 = 2
artist_dict['%s_m1'%event].set_zorder(zorder_1)
artist_dict['%s_m2'%event].set_zorder(zorder_2)
return artist_dict
anim = animation.FuncAnimation(fig,animate,init_func=init,frames=n_frames)
anim.save('{}/{}_orbit.gif'.format(args.outfolder,event),fps=fps,writer='imagemagick',extra_args=['-vcodec', 'libx264'])
if __name__ == '__main__':
main()
......@@ -12,6 +12,7 @@ setup(
description = __description__,
scripts = [
'bin/animate_binaries_from_file',
'bin/make_event_gifs',
],
packages = [
'twirl',
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment