-
Notifications
You must be signed in to change notification settings - Fork 1
/
mcfost_models.py
executable file
·549 lines (419 loc) · 15.5 KB
/
mcfost_models.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
#!/usr/bin/python3
from matplotlib.patches import Circle
from modules.casa_cube import casa_cube as casa
import matplotlib.pyplot as plt
import numpy as np
from modules.pymcfost import pymcfost as mcfost
from modules.colorbar_utils import colorbar2, shift_axes
from modules.params import Params
#------------------------------#
# Variables
#------------------------------#
# Set up the parameters
params = Params()
#------------------------------#
# Directory Information
#------------------------------#
# Path to the DSHARP data
dir = params.get("dir")
# Name of the simulation
name = params.get("name")
# Model directories
mod_dir_basename = dir if name.replace(" ", "") == ""else dir + name + "/"
# Name of the output PDF file
filename = mod_dir_basename + params.get("file")
#------------------------------#
#------------------------------#
# Observational Variables
#------------------------------#
# Include Observation Data
include_observation = params.get("inc_obs")
# The name of the observation
obs_name = params.get("obs_name")
# Directory to Observational Data
# Velocity CO channel must be called lines.fits.gz or lines.fits
# Continuum data must be called RT.fits.gz or RT.fits
dir_observation = dir + obs_name + "/"
# Match Observation Scales
# This will match the bmin, bmax and bpa from the observation
match_observation = params.get("match")
#------------------------------#
#------------------------------#
# System Variables
#------------------------------#
# Star velocity factor moving away
v_system = params.get("v_sys")
# Location of planet
p_loc = params.get_array("p_loc")
# Location of star shift
s_shift = [0, 0]
#------------------------------#
#------------------------------#
# Model Variables
#------------------------------#
# Flag for including models
include_models = params.get("inc_mods")
# The model names
models = params.get_array("models", strings = True) if include_models else []
# The mass of the planets (in Jupiter masses)
# For models without planets, write 0
# MAKE SURE THE SIZE OF THIS ARRAY IS THE SAME AS MODELS
p_masses = params.get_array("p_mass") if include_models else []
# Whether to use a moment, get the moment number
use_moment = params.get("use_moment")
moment = params.get("moment") if use_moment else 0
# Velocity channels
v_channels = params.get_array("v_chan") if not use_moment else [0]
#------------------------------#
#------------------------------#
# Plotting Variables
#------------------------------#
# Whether to plot continuum images or not
include_continuum = params.get("inc_cont")
# Whether to plot the channels or not
include_channels = params.get("inc_chan")
# The colour mapping for the observational plot
cmap_cont = "gist_earth"
# Whether to use Flux Temperature or not
c_plot_temp = True
# Whether to plot the location of sinks on the data or not
plot_sinks = params.get("plot_sinks")
# The minimum and maximum flux value for the pixels for velocity images
# These will scale the velocity channels
v_f_min = params.get("v_min")
v_f_max = params.get("v_max")
# The f_min value for the pixels in the continuum image
# The minimum and maximum flux value for the pixels on the continuum images
# These will scale the continuum
c_f_min = params.get("c_min")
c_f_max = params.get("c_max")
# The continuum pixel addition
# This makes the continuum darker the larger the value is
# Keep this as 0 if you do not want to change the scaling
c_mod_pix_add = 0
# The continuum colour scale
# Options are 'log' or 'lin'
c_color_scale = params.get("c_scale")
# The limits of the graph
# (max_x, min_x, min_y, max_y)
limits_max = params.get("limits")
limits = [limits_max,-limits_max,-limits_max,limits_max]
# Python Figure Size
f_size = 2.0
# Python spacing between figures
f_spacing = 0.12
#------------------------------#
#------------------------------#
# Applicaion Variables
# DO NOT CHANGE
#------------------------------#
# The number of models being used
n_models = len(models)
# The number of channels being used
n_channels = len(v_channels)
#------------------------------#
#------------------------------#
# Set up the Plots
#------------------------------#
# If including the observational data
if include_observation:
# Create continuum and CO data
if include_continuum:
cont = casa.Cube(dir_observation + "RT.fits")
if include_channels:
CO = casa.Cube(dir_observation + "lines.fits")
# Create the subplots
if include_channels:
fig, axes = plt.subplots(
nrows = n_models + int(include_observation),
ncols = (int(include_channels) * n_channels) + int(include_continuum),
figsize = (f_size * ((int(include_channels) * n_channels) + int(include_continuum)), f_size * (n_models + int(include_observation))),
sharex='all',
sharey='all'
)
# Otherwise make plots horizontal
else:
fig, axes = plt.subplots(
nrows = 1,
ncols = n_models + int(include_observation),
figsize = (f_size * (n_models + int(include_observation)), f_size * 1),
sharex='all',
sharey='all'
)
# Add some whitespace between them
plt.subplots_adjust(wspace = f_spacing, hspace = f_spacing)
# Adjust the axes for each row
if include_continuum:
for i in range(n_models + int(include_observation)):
try:
shift_axes(axes[i,0], -0.03, 0)
#shift_axes(axes[i,4:],0.01,0)
except:
pass
#------------------------------#
#------------------------------#
# Creates a circle at the planet position on the graph
def CreateCircle ():
return Circle(
(p_loc[0], p_loc[1]),
0.15,
clip_on = False,
zorder = 10,
linewidth = 2,
edgecolor = 'white',
linestyle = ":",
facecolor = (0, 0, 0, .0125),
alpha = 0.3,
fill = False
)
#------------------------------#
#------------------------------#
# Broadcast any Errors
#------------------------------#
# Check for mismatched arrays
if len(p_masses) < len(models):
raise Exception("Incorect Planet Masses Array Size. Make sure the array length is identical to the model array length.")
#------------------------------#
# Plot the Observational data
#------------------------------#
# If including the observational data
if include_observation:
# If plotting continuum graphs
if include_continuum:
# Get axis
axis = axes[0, 0] if include_channels and include_models else axes[0]
print("Plotting Observation Continuum")
# We plot the observations on the first row
image = cont.plot(
colorbar = False,
cmap = cmap_cont,
color_scale = c_color_scale,
ax = axis,
no_xlabel = include_channels,
no_ylabel = False,
limits = limits,
shift_dx = s_shift[0],
shift_dy = s_shift[1],
Tb = c_plot_temp,
fmin = c_f_min,
fmax = c_f_max,
)
# Show the colour bar in the plot
if include_channels or n_models <= 1:
colorbar2(image)
# Add the planet and star to the plot
axis.plot(s_shift[0], s_shift[1], "*", color="white", ms=4)
axis.plot(p_loc[0], p_loc[1], "o", color="cyan", ms=2)
# Label the planet name on the continuum plot
axis.text(
0.05,
0.9,
obs_name,
horizontalalignment = 'left',
color = "white",
transform = axis.transAxes,
fontsize = 10
)
# If including velocity channels
if include_channels:
print("Plotting Observation Velocity Channels")
# Loop though all the channels
for i in range(n_channels):
# Determine the current velocity for the observational data
iv = np.abs(CO.velocity - (v_system + v_channels[i])).argmin()
# Only show color bar in the last channel
show_colorbar = i == n_channels - 1
# Get the axis
if len(axes) > 0 and type(axes[0]) is np.ndarray:
ax = axes[0, int(include_continuum) + i] if include_models else axes[int(include_continuum) + i]
else:
ax = axes[int(include_continuum) + i]
# Plot the velocity channel
vel_im = CO.plot(
iv = iv,
v0 = v_system,
colorbar = False,
ax = ax,
no_xlabel = True,
no_ylabel = True,
limits = limits,
shift_dx = s_shift[0],
shift_dy = s_shift[1],
Tb = c_plot_temp,
fmax = v_f_max if not use_moment else None,
fmin = v_f_min if not use_moment else None,
moment = moment if use_moment else None
)
if show_colorbar:
colorbar2(vel_im)
# Add a circle where the planet is expected to be
if plot_sinks:
circle = CreateCircle()
ax.add_artist(circle)
# If no previous label
if i == 0 and not include_continuum:
ax.text(
0.05,
0.9,
obs_name,
horizontalalignment = 'left',
color = "white",
transform = ax.transAxes,
fontsize = 10
)
#------------------------------#
#------------------------------#
# Plot the Simulation Models
#------------------------------#
# Loop through each of the simulaions
for k, mod in enumerate(models):
# Get the model directory
mod_dir = mod_dir_basename + str(mod)
# Print message
print("Analysing output from {} located at:\n\t{}".format(mod, mod_dir))
# Get the continuum and CO images
if include_continuum: mod_cont = mcfost.Image(mod_dir)
mod_CO = mcfost.Line(mod_dir)
# Determine whether to display the xlabel or not
no_xlabel = k < n_models - 1
# Add pixel values to the continuum image
if include_continuum: mod_cont.image += c_mod_pix_add * mod_cont.image[4,0,0,:,:]
# If plotting the continuum
if include_continuum:
# Get the axis
if len(axes) > 0 and type(axes[0]) is np.ndarray:
axis = axes[k + int(include_observation), 0] if include_channels else axes[k + int(include_observation)]
else:
ax = axes[k + int(include_observation)]
# Plot the continuum
if include_observation and match_observation:
image = mod_cont.plot(
ax = axis,
colorbar = False,
bmaj = cont.bmaj,
bmin = cont.bmin,
bpa = cont.bpa,
no_xlabel = no_xlabel and include_channels,
no_ylabel = not include_channels,
limits = limits,
cmap = cmap_cont,
scale = c_color_scale,
Tb = c_plot_temp,
vmin = c_f_min,
vmax = c_f_max,
plot_stars = plot_sinks,
)
# If no observational data to base scales on
else:
image = mod_cont.plot(
ax = axis,
colorbar = False,
no_xlabel = no_xlabel and include_channels,
no_ylabel = not include_channels,
limits = limits,
cmap = cmap_cont,
scale = c_color_scale,
Tb = c_plot_temp,
vmin = c_f_min,
vmax = c_f_max,
plot_stars = plot_sinks,
)
# Label the model on the continuum plot
if p_masses[k] == 0:
mod_label = models[k] + " Model"
else:
mod_label = str(p_masses[k]) + ("$\mathrm{M_J}$ Model")
axis.text(
0.05,
0.9,
mod_label,
horizontalalignment = 'left',
color = "white",
transform = axis.transAxes,
fontsize = 10
)
# Add the planet and star to the plot
#axes[k+int(include_observation),0].plot(s_shift[0], s_shift[1], "*", color="white", ms=4)
#if mplanet != 0:
# axes[k+int(include_observation),0].plot(p_loc[0], p_loc[1], "o", color="cyan", ms=2)
# Show the colour bar in the plot
if include_channels or k == n_models - 1:
colorbar2(image)
# Set the change in velocity
delta_v = None
# If including velocity channels
if include_channels:
# Loop through each of the channels
for i in range(n_channels):
# Get the axis
if len(axes) > 0 and type(axes[0]) is np.ndarray:
axis = axes[k + int(include_observation), int(include_continuum) + i]
else:
axis = axes[k + int(include_observation)]
# Only show color bar in the last channel
show_colorbar = i == n_channels - 1
# Plot the CO velocity channel
if include_observation and match_observation:
vel_im = mod_CO.plot_map(
v = v_channels[i],
ax = axis,
colorbar = False,
bmaj = CO.bmaj,
bmin = CO.bmin,
bpa = CO.bpa,
no_xlabel = no_xlabel,
no_ylabel = True,
limits = limits,
Tb = c_plot_temp,
Delta_v = delta_v,
fmax = v_f_max if not use_moment else None,
fmin = v_f_min if not use_moment else None,
moment = moment if use_moment else None,
plot_stars = plot_sinks,
vlabel_size = 8,
)
# If no observational data to base scales on
else:
vel_im = mod_CO.plot_map(
v = v_channels[i],
ax = axis,
colorbar = False,
no_xlabel = no_xlabel,
no_ylabel = True,
limits = limits,
Tb = c_plot_temp,
Delta_v = delta_v,
fmax = v_f_max if not use_moment else None,
fmin = v_f_min if not use_moment else None,
moment = moment if use_moment else None,
plot_stars = plot_sinks,
vlabel_size = 8,
)
# Plot the circle where the planet is expected to be
if p_masses[k] != 0 and plot_sinks:
circle = CreateCircle()
axis.add_artist(circle)
# Add in a colorbar
if show_colorbar:
colorbar2(vel_im)
# If no label has been given for the other models
if not include_continuum and i == 0:
# Label the model on the continuum plot
if p_masses[k] == 0:
mod_label = models[k] + " Model"
else:
mod_label = str(p_masses[k]) + ("$_\mathrm{M_J}$ Model")
axis.text(
0.05,
0.9,
mod_label,
horizontalalignment = 'left',
color = "white",
transform = axis.transAxes,
fontsize = 10
)
#------------------------------#
# Save the figure
plt.savefig(filename, bbox_inches='tight')
# Show the graph in an xw display window
plt.show()