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

updated animate_binaries_from_file to take in multiple catalogs and special names

parent 5fa8e63b
......@@ -13,9 +13,10 @@ def make_parser():
parser = ArgumentParser()
parser.add_argument(
"file",
"files",
type=str,
help='json file with median values or gwosc text catalog',
nargs='*'
)
parser.add_argument(
......@@ -43,9 +44,36 @@ def make_parser():
default=False
)
parser.add_argument(
"--name_file",
help="file that maps to preferred event names",
type=str,
)
parser.add_argument(
"--nxny",
help="number of x and y subplots. must be equal to number of events, two arguments",
type=int,
nargs=2
)
parser.add_argument(
"--figsize",
help="x and y figure size, 2 arguments",
type=float,
nargs=2
)
parser.add_argument(
"--object_scale",
help="scale of objects drawn",
default = 1./200,
type = float
)
return parser
def file_handler(filename):
def file_handler(filename,names=None):
""" spits out a dictionary of events mapping
to m1 and m2 values
"""
......@@ -56,7 +84,11 @@ def file_handler(filename):
data = json.load(json_file)
for event in list(data.keys()):
events2masses[event] = {
if names:
ev_name = names[event]
else:
ev_name = event
events2masses[ev_name] = {
'm1': data[event]['mass_1_source']['med'],
'm2': data[event]['mass_2_source']['med']
}
......@@ -68,7 +100,8 @@ def file_handler(filename):
for event in data['id']:
print(event)
event_sel = data['id'] == event
events2masses[event] = {
event_common_name = data['commonName'][np.argwhere(event_sel)]
events2masses[event_common_name[0][0]] = {
'm1': float(data[event_sel]['mass_1_source'][0]),
'm2': float(data[event_sel]['mass_2_source'][0])
}
......@@ -83,21 +116,40 @@ def main():
# parse arguments
parser = make_parser()
args = parser.parse_args()
# get the names if needed:
if args.name_file:
with open(args.name_file,'r') as filey:
names = json.load(filey)['NNAME']
else:
names = None
events2masses = file_handler(args.file)
events2masses = {}
for filey in args.files:
events2masses.update(file_handler(filey,names=names))
# get number of events and hence number of axes needed
events = list(events2masses.keys())
n_events = len(events)
n_side = int(np.ceil(np.sqrt(n_events)))
n_side2 = int(np.ceil(n_events/float(n_side)))
if args.nxny:
n_side,n_side2 = args.nxny
else:
n_side = int(np.ceil(np.sqrt(n_events)))
n_side2 = int(np.ceil(n_events/float(n_side)))
# make axes
if args.dark:
plt.style.use('dark_background')
fig,ax = plt.subplots(n_side,n_side2,figsize=(10,10))
if args.figsize:
figx,figy = args.figsize
else:
figx,figy = 10,10
fig,ax = plt.subplots(n_side,n_side2,figsize=(figx,figy),gridspec_kw = {'hspace':0.25})
ax = ax.flatten()
for axi in ax: axi.set_axis_off()
for axi in ax:
axi.set_axis_off()
axi.axis("equal")
# set frame rate and movie length
fps = 30. # frames per second
......@@ -119,7 +171,7 @@ def main():
pos_dict = {} # dictionary for all compact object positions
bbh_dict = {} # dictionary for all BinaryBlackHole objects
# loop over all events
for ei,event in enumerate(events):
for ei,event in enumerate(sorted(events)):
print('evolving ',event)
event_pos_dict = {} # dictionary for positions of this event
......@@ -146,18 +198,20 @@ def main():
binary = binary,
ax = ax[ei],
name=event,
BH_scale = 1./200
BH_scale = args.object_scale
)
bbh_dict[event] = bbh
artist_dict = {} # dictionary for matplotlib artists
def init():
for ei,event in enumerate(events):
for ei,event in enumerate(sorted(events)):
artist_dict.update(bbh_dict[event].setup_artists())
# configure plot axes, title
ax[ei].set_xlim([-1,1])
ax[ei].set_ylim([-1,1])
#ax[ei].set_box_aspect(1)
ax[ei].set_aspect('equal')
ax[ei].set_xlim([-1.3,1.3])
ax[ei].set_ylim([-1.3,1.3])
print(args.no_names)
if not args.no_names:
print('putting names on events')
......@@ -165,13 +219,15 @@ def main():
return artist_dict
def animate(i):
for ei,event in enumerate(events):
for ei,event in enumerate(sorted(events)):
artist_dict.update(bbh_dict[event].get_artists(i))
#ax[ei].set_box_aspect(1)
return artist_dict
anim = animation.FuncAnimation(fig,animate,init_func=init,frames=n_frames)
#anim.save('orbit.gif',fps=fps,writer='imagemagick',extra_args=['-vcodec', 'libx264'])
anim.save(args.outfile,fps=fps,extra_args=['-vcodec', 'libx264'])
#anim.save(f'{args.outfile}.gif',fps=fps,writer='imagemagick',extra_args=['-vcodec', 'libx264'])
anim.save(f'{args.outfile}.mp4',fps=fps,extra_args=['-vcodec', 'libx264'])
if __name__ == '__main__':
main()
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