Create absorption from `AthenaK` data using `absorption_line_profile_calculator.py`

absorption_tests.py

# -*- coding: utf-8 -*-
"""
Created on Fri Nov  17 15:20:28 2023

@author: alankar
"""

import sys
import numpy as np
import yt
from yt import derived_field
import matplotlib.pyplot as plt
import astropy.constants as const
import absorption_line_profile_calculator as absorption
# import kit_np as ak
# plt.style.use('dark_background')

kpc = const.kpc.cgs.value # in cm
kB = const.k_B.cgs.value
mp = const.m_p.cgs.value
amu = const.u.cgs.value
c = const.c.cgs.value
XH  = 0.7154
YHe = 0.2703
ZMet = 0.0143

mu = 0.6000317       # mean molecular weight
gamma = 5./3.

prefix='Turb'
varnames='.hydro_w.'
bin_suffix='.bin'
hdf_suffix='.athdf'
xdmf_suffix = '.athdf.xdmf'
parentdir = '/home/alankar/Downloads/'
filedir = 'with_cool_multiphase/bin/'
filedir = 'with_cool_512/bin/'
filedir = ''
filenumber=20

hdf_filename=parentdir+filedir+prefix+varnames+str(filenumber).rjust(5, '0')+hdf_suffix

length_code = kpc
time_code = 3.15576e+15 # 100 Myr
mass_code = 3.036951775493658e+39

# Prepare and load Athena data
@derived_field(name="temp", units="code_length**3/code_mass", sampling_type="cell")
def _temp(field, data):
    return (gamma-1)*data["athena_pp", "eint"]/data["athena_pp", "dens"]*mu*amu/kB

units_override = {
    "length_unit": (length_code, "cm"),
    "time_unit": (time_code, "s"),
    "mass_unit": (mass_code, "g"),
}
parameters = {
    "gamma": gamma,
    "geometry": "cartesian",
    "periodicity": (True, True, True),
}
ds = yt.load(hdf_filename, units_override=units_override, parameters=parameters)

# Generate a ray
ray = ds.r[0.0,-0.37,-0.5:0.5]
dens = np.array(ray[("athena_pp","dens")])*mass_code/length_code**3
nH = dens*XH/amu
Temp = np.array(ray[("gas","temp")])
dx = np.array(ray[("athena_pp","dx")])*length_code
velx = np.array(ray["athena_pp","velx"])*length_code/time_code 

load = True
n_MgII = np.zeros_like(nH)
if not(load):
    sys.path.append('/home/alankar/Documents/MultiphaseGalacticHaloModel/submodules/AstroPlasma/')
    from astro_plasma import Ionization
    fIon = Ionization.interpolate_ion_frac
    num_dens = Ionization.interpolate_num_dens

    # Prepare plasma condition for AstroPlasma
    metallicity = 0.33
    redshift = 0.001
    mode="PIE"
    element = "MgII"

    for i in range(nH.shape[0]):
        print(i, end="\r")
        temperature = Temp[i]
        vlos = velx[i]
        nH = nH[i] 
        n_MgII[i] = num_dens(nH = nH,
                  temperature = temperature,
                  metallicity = metallicity,
                  redshift = redshift,
                  mode = mode,
                  element = element)
    np.save('athena_MgII.npy', n_MgII)
else:
    n_MgII = np.load("athena_MgII.npy")
    
# Plotting Athena data along the ray
z = np.array(ray[("athena_pp","z")])
plt.figure()
plt.plot(z, nH, label='n_H')
plt.plot(z,n_MgII, label='n_MgII')
plt.legend()
plt.yscale('log')
plt.xlim(-0.5, 0.5)
plt.xlabel('z')
plt.ylabel('n')
plt.ylim(ymin=1e-12)
plt.grid()
plt.savefig('n_ion_nH.png')
plt.show()

dz = np.array(ray[("athena_pp","dz")])*length_code
velz = np.array(ray["athena_pp","velz"])*length_code/time_code
cs = np.sqrt(gamma*kB*Temp/(mu*mp))

plt.figure()
plt.plot(z, velz/cs, label='Mach_z')
plt.legend()
plt.xlim(-0.5, 0.5)
plt.xlabel('z')
plt.ylabel('Mach_z_loc')
plt.grid()
plt.savefig('Mach_z_loc.png')
plt.show()

plt.figure()
plt.plot(z, velz/1e5, label='vel_z')
plt.legend()
plt.xlim(-0.5, 0.5)
plt.xlabel('z')
plt.ylabel('vel_z (km/s)')
plt.grid()
plt.savefig('vel_z.png')
plt.show()

# Mg II atomic data for transition
a_ion = 24.305 # Mg atomic mass in amu
wave0 = 2796.3553 # rest transition wavelength in angstrom
f_lu = 6.155e-01 # oscillator strength in CGS
gamma_ul = 2.625e+08 # sum of einstein coefficients, i.e., total transition probability

vel_range = 2.0*np.max(np.abs(velz))/1.0e+05 # km/s
spec_res = [0.01, 2.2] # km/s
tau_eff = None
freq = None
column = 0

for i in range(nH.shape[0]):
    temperature = Temp[i]
    del_x = dz[i]
    n_ion = n_MgII[i]
    vlos = velz[i]/1.0e+05 # km/s

    # only 1 transition
    if tau_eff is None:
        freq, tau_eff = absorption.generate_absorption(vel_range, spec_res[0], temperature, vlos, n_ion*del_x, 
                        a_ion, wave0, f_lu, gamma_ul)
    else:
        tau_eff += absorption.generate_absorption(vel_range, spec_res[0], temperature, vlos, n_ion*del_x, 
                        a_ion, wave0, f_lu, gamma_ul)[1]
    column += (n_ion*del_x)
    
print("MgII Column %e cm^-2"%column)

Flux = absorption.generate_norm_flux(tau_eff)
del_wave_coarse  = wave0/(c/1.0e+05/spec_res[1])
wave_angstrom = (c/freq)*1.0e+08
wave_angstrom_lowres, Flux_lowres = absorption.calc_flux_coarse(wave_angstrom, 
                                                                Flux, 
                                                                del_wave_coarse)

# Plot Flux

# wavellength space
fig = plt.figure()
plt.plot(wave_angstrom, Flux)
absorption.plot_absorption(wave_angstrom_lowres,  Flux_lowres, fig, plt.gca())
# plt.xscale('log')
# plt.yscale('log')
plt.xlim(np.min(wave_angstrom), np.max(wave_angstrom))
plt.xlabel('lambda')
plt.ylabel('F/F_0')
plt.grid()
plt.savefig('flux.png')
plt.show()

# LOS velocity space
freq_lowres = c/(wave_angstrom_lowres*1.0e-08)
wave_vel_lowres = absorption.generate_velocity_base(freq_lowres, wave0)
wave_vel = absorption.generate_velocity_base(freq, wave0)
fig = plt.figure()
plt.plot(wave_vel, Flux)
absorption.plot_absorption(wave_vel_lowres, Flux_lowres, fig, plt.gca())
# plt.xscale('log')
# plt.yscale('log')
plt.xlim(-vel_range, vel_range)
plt.xlabel('vel_spect (km/s)')
plt.ylabel('F/F_0')
plt.grid()
plt.savefig('flux_vel_spect.png')
plt.show()

EW = absorption.calculate_EW(wave_angstrom_lowres, Flux_lowres) # mÅ
print("MgII 2796 EW (LoRes) = %.1f mÅ"%EW)
EW = absorption.calculate_EW(wave_angstrom, Flux) # mÅ
print("MgII 2796 EW (HiRes) = %.1f mÅ"%EW)