plot_avir_sim.py

"""
Script to generate synthetic virial plots from the Collins+ 2012
simulation data set taken from CATS. This script generates figures 3
and 4 of Krumholz, Lada, & Forbrich (2025).
"""

import h5py
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colormaps as cm
import astropy.units as u
from avir_sim_tools import project_quantities, \
    contour_analysis, contour_summary, compute_LCO

# Simulation data files at various times
fname = ['../C12_Beta0.2_256_0000.h5',
         '../C12_Beta0.2_256_0010.h5',
         '../C12_Beta0.2_256_0030.h5',
         '../C12_Beta0.2_256_0060.h5']
times = np.array([ 0, 0.1, 0.3, 0.6 ])

# Set contour levels
clev = np.linspace(0.1, 0.7, 25)

# Set minimum number of pixels to keep; 30 corresponds to a radius of
# ~= 3 pixels
min_pix = 30

# Numerical value of G in simulation unit system where mean box
# density = 1, box size = 1, sound speed = 1
G = 5 * 9**2 / 3

# Read data
data = []
for f in fname:
    fptr = h5py.File(f, 'r')
    data.append( {
        'rho' : np.array(fptr['density']),
        'vx' : np.array(fptr['velocity_x']),
        'vy' : np.array(fptr['velocity_y']),
        'vz' : np.array(fptr['velocity_z'])
        } )
    fptr.close()

# Compute 12CO and 13CO luminosities
for d in data:
    compute_LCO(d, nHscale = 1e2*u.cm**-3)

# Compute projections
for d in data:
    project_quantities(d)

# Do contour analysis on data
for d in data:
    d['contour_data'] = contour_analysis(d, clev, G, min_pix = min_pix)
    
# Compute summary statistics
for d in data:
    d['contour_summary'] = contour_summary(d['contour_data'],
                                           group_proj=True,
                                           ncmin=5)    

# Set fonts for plots
plt.rc('text', usetex=True)
plt.rc('font', family='serif', size=12)

# Create plot window for first figure
fig = plt.figure(1, figsize=(3.5, 5.5))
fig.clf()

# Set limits
avir_lim = np.array([-0.3, 0.5])
cden_lim = np.array([0.05, 0.7])
s2r_lim = np.array([1.8, 2.8])

# Grab line colors
ntime = len(times)
colors = [ cm['copper'](i / (ntime-1)) for i in range(ntime) ]

# Top panel: avir_obs
ax = fig.add_subplot(2,1,1)
for i in range(ntime):
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['avir_mean']),
            color=colors[i],
            label="{:3.1f}".format(times[i]))
ax.plot(cden_lim, np.log10(1)*np.ones(2), 'k--', lw=1, alpha=0.3)
ax.plot(cden_lim, np.log10(2)*np.ones(2), 'k--', lw=1, alpha=0.3)
ax.text(0.5, np.log10(1), r'$\alpha_\mathrm{vir,obs} = 1$',
        fontdict = {'size' : 10}, va='bottom')
ax.text(0.5, np.log10(2), r'$\alpha_\mathrm{vir,obs} = 2$',
        fontdict = {'size' : 10}, va='bottom')
ax.set_xlim(cden_lim)
ax.set_ylim(avir_lim)
ax.set_xticklabels('')
ax.set_ylabel(r'$\log\alpha_\mathrm{vir,obs}$')

# Bottom panel: sigma^2 / r
ax = fig.add_subplot(2,1,2)
for i in range(len(times)):
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['s2r_mean']),
            color=colors[i],
            label="{:3.1f}".format(times[i]))
ax.plot(cden_lim,
        np.log10(1*G*10.**cden_lim*np.pi/5), 'k--', lw=1, alpha=0.3)
ax.plot(cden_lim,
        np.log10(2*G*10.**cden_lim*np.pi/5), 'k--', lw=1, alpha=0.3)
ax.text(0.5, np.log10(1*G*10.**0.5*np.pi/5)+0.13,
        r'$\alpha_\mathrm{vir,obs} = 1$',
        rotation=29,
        fontdict = {'size' : 10}, va='top')
ax.text(0.1, np.log10(2*G*10.**0.1*np.pi/5)+0.02,
        r'$\alpha_\mathrm{vir,obs} = 2$',
        rotation=29,
        fontdict = {'size' : 10}, va='bottom')
ax.set_xlim(cden_lim)
ax.set_ylim(s2r_lim)
ax.legend(title=r'$t/t_\mathrm{ff}$', title_fontsize=10,
          prop={'size' : 10},
          ncol=2, loc='lower right')
ax.set_xlabel(r'$\log\Sigma$')
ax.set_ylabel(r'$\log\langle\sigma^2\rangle/R_\mathrm{eff}$')

# Adjust spacing and save
plt.subplots_adjust(top=0.95, hspace=0.05, right=0.95, left=0.2)
plt.savefig('avir_sim.pdf')


# Make a second plot illustrating the procedure
fig = plt.figure(2, figsize=(3.5,6.25))
fig.clf()

# Log column density limits
cdmin = -1
cdmax = 1

# Contour level to illustrate
clevex = 0.25

# Create upper image
ax1 = fig.add_subplot(2,1,1)
im = ax1.imshow(np.log10(data[1]['cden']['z']).T, origin='lower',
                aspect='equal', extent=(0,1,0,1),
                vmin=cdmin, vmax=cdmax,
                cmap='coolwarm')
ax1.contour(np.log10(data[1]['cden']['z']).T, levels=[clevex],
            extent=(0,1,0,1), linewidths=[0.5])
ax1.text(0.65, 0.9, r'$t/t_\mathrm{ff} = 0.1$',
         bbox={'facecolor' : 'white', 'alpha' : 0.5 })
ax1.set_xticklabels([])
ax1.set_ylabel(r'$y$')

# Create lower image
ax2 = fig.add_subplot(2,1,2)
ax2.imshow(np.log10(data[3]['cden']['z']).T, origin='lower',
           aspect='equal', extent=(0,1,0,1),
           vmin=cdmin, vmax=cdmax, cmap='coolwarm')
ax2.contour(np.log10(data[3]['cden']['z']).T, levels=[clevex],
            extent=(0,1,0,1), linewidths=[0.5])
ax2.text(0.65, 0.9, r'$t/t_\mathrm{ff} = 0.6$',
         bbox={'facecolor' : 'white', 'alpha' : 0.5 })
ax2.set_ylabel(r'$y$')
ax2.set_xlabel(r'$x$')

# Adjust spacing
fig.subplots_adjust(left=0.1, top=0.88, right=1, hspace=0.07)

# Add color bar
fig.draw_without_rendering()  # Needed to force coordinates to update
coord1 = fig.transFigure.inverted().transform(
    ax1.transData.transform([0,1.03]))
coord2 = fig.transFigure.inverted().transform(
    ax1.transData.transform([1,1.08]))
axc = fig.add_axes((coord1[0], coord1[1],
                    coord2[0]-coord1[0],
                    coord2[1]-coord1[1]))
fig.colorbar(im, cax=axc, orientation='horizontal')
axc.xaxis.set_ticks_position('top')
axc.xaxis.set_label_position('top')
axc.set_xlabel(r'$\log\Sigma$')

# Save
plt.savefig('sim_contour.pdf')


# Make a third plot showing histograms of the ratio of 12CO and 13CO
# column density and velocity dispersion to true values
fig = plt.figure(3, figsize=(3.5, 5.0))
fig.clf()
xlim = [-0.5, 0.75]
ylim = [0,8]

# First column: column densities

# Loop over direction and time to aggregate ratios
ratio12 = []
ratio13 = []
cden = []
for d in data:
    ratio12 = ratio12 + [np.array([])]
    ratio13 = ratio13 + [np.array([])]
    for d1 in ['x', 'y', 'z']:
        ratio12[-1] = np.append(ratio12[-1],
                                np.sqrt( d['cden_12'][d1] /
                                         d['cden'][d1] ) )
        ratio13[-1] = np.append(ratio13[-1],
                                np.sqrt( d['cden_13'][d1] /
                                         d['cden'][d1] ) )

# Plot histogram of ratios
for i in range(len(times)):
    ax = fig.add_subplot(len(times),2,1+2*i)
    ax.hist(np.log10(ratio12[i]), bins=40, range=xlim, density=True,
            color='C0', alpha=0.5, label=r'$^{12}$CO')
    ax.hist(np.log10(ratio13[i]), bins=40, range=xlim, density=True,
            color='C1', alpha=0.5, label=r'$^{13}$CO')
    ax.plot(np.log10(np.median(ratio12[i])) * np.ones(2),
            ylim, 'C0--')
    ax.plot(np.log10(np.median(ratio13[i])) * np.ones(2),
            ylim, 'C1--')
    ax.text(xlim[1]-0.03, 7.5,
            r'$t/t_\mathrm{{ff}} = {:3.1f}$'.format(times[i]),
            ha='right', va='top', fontsize=10)
    if i == len(times)-1:
        ax.set_xlabel(r'$\log \Sigma_\mathcal{L}/\Sigma$')
    else:
        ax.set_xticks([])
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    ax.set_ylabel('PDF')

# Second column: velocity dispersions

# Loop over direction and time to aggregate ratios
ratio12 = []
ratio13 = []
cden = []
for d in data:
    ratio12 = ratio12 + [np.array([])]
    ratio13 = ratio13 + [np.array([])]
    for d1 in ['x', 'y', 'z']:
        ratio12[-1] = np.append(ratio12[-1],
                                np.sqrt( d['sigma2_12'][d1] /
                                         d['sigma2'][d1] ) )
        ratio13[-1] = np.append(ratio13[-1],
                                np.sqrt( d['sigma2_13'][d1] /
                                         d['sigma2'][d1] ) )

# Plot histogram of ratios
ylim = [0,7.3]
for i in range(len(times)):
    ax = fig.add_subplot(len(times),2,2+2*i)
    ax.hist(np.log10(ratio12[i]), bins=40, range=xlim, density=True,
            color='C0', alpha=0.5, label=r'$^{12}$CO')
    ax.hist(np.log10(ratio13[i]), bins=40, range=xlim, density=True,
            color='C1', alpha=0.5, label=r'$^{13}$CO')
    ax.plot(np.log10(np.median(ratio12[i])) * np.ones(2),
            ylim, 'C0--')
    ax.plot(np.log10(np.median(ratio13[i])) * np.ones(2),
            ylim, 'C1--')
    if i == 0:
        ax.legend(prop={'size' : 8}, loc='upper right')
    if i == len(times)-1:
        ax.set_xlabel(r'$\log \sigma_\mathcal{L}/\sigma$')
    else:
        ax.set_xticks([])
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    ax.set_yticks([])

fig.subplots_adjust(left=0.15, bottom=0.1, top=0.95, right=0.95,
                    hspace=0.05, wspace=0.05)

# Save
plt.savefig('sigma12_13_ratio.pdf')

        

# Create plot window for fourth figure
fig = plt.figure(4, figsize=(6.5, 4))
fig.clf()

# Top left panel: avir_obs for line+dust
ax = fig.add_subplot(2,2,1)

for i in range(ntime):
    # Labels
    if i == 0:
        label_true = 'True'
        label_12 = r'$^{12}$CO'
        label_13 = r'$^{13}$CO'
    else:
        label_true = None
        label_12 = None
        label_13 = None
        
    # Means
    mean = np.nanmean(data[i]['contour_summary']['avir_mean'])
    mean12m = np.nanmean(data[i]['contour_summary']['avir_12m_mean'])
    mean13m = np.nanmean(data[i]['contour_summary']['avir_13m_mean'])

    # Plot true, 12CO, and 13CO data data
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['avir_mean']),
            color=colors[i],
            label=label_true)
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['avir_12m_mean']) -
            np.log10(mean12m) + np.log10(mean),
            ls='--',
            color=colors[i],
            label=label_12)
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['avir_13m_mean']) -
            np.log10(mean13m) + np.log10(mean),
            ls=':',
            color=colors[i],
            label=label_13)
ax.plot(cden_lim, np.log10(1)*np.ones(2), 'k--', lw=1, alpha=0.3)
ax.plot(cden_lim, np.log10(2)*np.ones(2), 'k--', lw=1, alpha=0.3)
ax.text(0.5, np.log10(1), r'$\alpha_\mathrm{vir,obs} = 1$',
        fontdict = {'size' : 10}, va='bottom')
ax.text(0.5, np.log10(2), r'$\alpha_\mathrm{vir,obs} = 2$',
        fontdict = {'size' : 8}, va='bottom')
ax.legend(prop={'size' : 8}, loc='upper left')
ax.set_title('Line + dust')
ax.set_xlim(cden_lim)
ax.set_ylim(avir_lim)
ax.set_xticklabels('')
ax.set_ylabel(r'$\log\alpha_\mathrm{vir,obs}$')


# Top right panel: avir_obs for line only
ax = fig.add_subplot(2,2,2)

for i in range(ntime):
    # Labels
    if i == 0:
        label_true = 'True'
        label_12 = r'$^{12}$CO'
        label_13 = r'$^{13}$CO'
    else:
        label_true = None
        label_12 = None
        label_13 = None
        
    # Means
    mean = np.nanmean(data[i]['contour_summary']['avir_mean'])
    mean12 = np.nanmean(data[i]['contour_summary']['avir_12_mean'])
    mean13 = np.nanmean(data[i]['contour_summary']['avir_13_mean'])

    # Plot true, 12CO, and 13CO data data
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['avir_mean']),
            color=colors[i],
            label=label_true)
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['avir_12_mean']) -
            np.log10(mean12) + np.log10(mean),
            ls='--',
            color=colors[i],
            label=label_12)
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['avir_13_mean']) -
            np.log10(mean13) + np.log10(mean),
            ls=':',
            color=colors[i],
            label=label_13)
ax.plot(cden_lim, np.log10(1)*np.ones(2), 'k--', lw=1, alpha=0.3)
ax.plot(cden_lim, np.log10(2)*np.ones(2), 'k--', lw=1, alpha=0.3)
ax.text(0.5, np.log10(1), r'$\alpha_\mathrm{vir,obs} = 1$',
        fontdict = {'size' : 10}, va='bottom')
ax.text(0.5, np.log10(2), r'$\alpha_\mathrm{vir,obs} = 2$',
        fontdict = {'size' : 10}, va='bottom')
ax.set_xlim(cden_lim)
ax.set_ylim(avir_lim)
ax.set_xticklabels('')
ax.set_yticklabels('')
ax.set_title('Line only')


# Bottom left panel: sigma^2 / r for line + dust
ax = fig.add_subplot(2,2,3)
for i in range(len(times)):

    # Means
    mean = np.nanmean(data[i]['contour_summary']['s2r_mean'])
    mean12 = np.nanmean(data[i]['contour_summary']['s2r_12m_mean'])
    mean13 = np.nanmean(data[i]['contour_summary']['s2r_13m_mean'])

    # Plot trure, 12CO, 13CO
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['s2r_mean']),
            color=colors[i],
            label="{:3.1f}".format(times[i]))
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['s2r_12m_mean']) -
            np.log10(mean12) + np.log10(mean),
            ls='--',
            color=colors[i])
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['s2r_13m_mean']) -
            np.log10(mean13) + np.log10(mean),
            ls=':',
            color=colors[i])

ax.plot(cden_lim,
        np.log10(1*G*10.**cden_lim*np.pi/5), 'k--', lw=1, alpha=0.3)
ax.plot(cden_lim,
        np.log10(2*G*10.**cden_lim*np.pi/5), 'k--', lw=1, alpha=0.3)
ax.text(0.5, np.log10(1*G*10.**0.5*np.pi/5)+0.13,
        r'$\alpha_\mathrm{vir,obs} = 1$',
        rotation=19,
        fontdict = {'size' : 10}, va='top')
ax.text(0.1, np.log10(2*G*10.**0.1*np.pi/5)+0.02,
        r'$\alpha_\mathrm{vir,obs} = 2$',
        rotation=19,
        fontdict = {'size' : 10}, va='bottom')
ax.set_xlim(cden_lim)
ax.set_ylim(s2r_lim)
ax.legend(title=r'$t/t_\mathrm{ff}$', title_fontsize=10,
          prop={'size' : 10},
          ncol=2, loc='lower right')
ax.set_xlabel(r'$\log\Sigma$')
ax.set_ylabel(r'$\log\langle\sigma^2\rangle/R_\mathrm{eff}$')


# Bottom right panel: sigma^2 / r for line only
ax = fig.add_subplot(2,2,4)
for i in range(len(times)):

    # Means
    mean = np.nanmean(data[i]['contour_summary']['s2r_mean'])
    mean12 = np.nanmean(data[i]['contour_summary']['s2r_12_mean'])
    mean13 = np.nanmean(data[i]['contour_summary']['s2r_13_mean'])

    # Plot trure, 12CO, 13CO
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['s2r_mean']),
            color=colors[i],
            label="{:3.1f}".format(times[i]))
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['s2r_12_mean']) -
            np.log10(mean12) + np.log10(mean),
            ls='--',
            color=colors[i])
    ax.plot(clev,
            np.log10(data[i]['contour_summary']['s2r_13_mean']) -
            np.log10(mean13) + np.log10(mean),
            ls=':',
            color=colors[i])

ax.plot(cden_lim,
        np.log10(1*G*10.**cden_lim*np.pi/5), 'k--', lw=1, alpha=0.3)
ax.plot(cden_lim,
        np.log10(2*G*10.**cden_lim*np.pi/5), 'k--', lw=1, alpha=0.3)
ax.text(0.5, np.log10(1*G*10.**0.5*np.pi/5)+0.13,
        r'$\alpha_\mathrm{vir,obs} = 1$',
        rotation=19,
        fontdict = {'size' : 10}, va='top')
ax.text(0.1, np.log10(2*G*10.**0.1*np.pi/5)+0.02,
        r'$\alpha_\mathrm{vir,obs} = 2$',
        rotation=19,
        fontdict = {'size' : 10}, va='bottom')
ax.set_xlim(cden_lim)
ax.set_ylim(s2r_lim)
ax.set_xlabel(r'$\log\Sigma$')
ax.set_yticklabels('')

# Adjust spacing and save
plt.subplots_adjust(top=0.9, hspace=0.05, right=0.95, left=0.12,
                    wspace=0.05, bottom=0.12)
plt.savefig('avir_sim_12_13.pdf')