Calculate absorption profiles

absorption_line_profile_calculator.py

# -*- coding: utf-8 -*-
"""
Created on Fri Nov  17 12:22:19 2023

@author: alankar
"""

import numpy as np
from scipy.special import voigt_profile
import astropy.constants as const

generate_norm_flux = lambda optical_depth: np.exp(-optical_depth)

def calculate_EW(
        wavel,      # wavelength in angstrom
        flux_norm,  # normalized flux F/F_0
        ): 
    # Draine eq 9.4
    return np.trapz((1.-flux_norm), wavel)*1.0e+03 # in milliangstrom

def calc_flux_coarse(
        wavel,      # wavelength in angstrom
        flux_norm,  # normalized flux F/F_0
        del_wave, # spectral resolution in angstrom
        ):
    nbins = int((np.max(wavel)-np.min(wavel))//del_wave)
    wavel_coarse = np.linspace(np.min(wavel), np.max(wavel), nbins+1)
    flux_coarse = np.zeros(nbins)
    for i in range(nbins):
        if i<nbins-1:
            condition = np.logical_and(wavel>=wavel_coarse[i], wavel<wavel_coarse[i+1])
        else:
            condition = np.logical_and(wavel>=wavel_coarse[i], wavel<=wavel_coarse[i+1])
        bin_width = np.max(wavel[condition]) - np.min(wavel[condition])
        flux_coarse[i] = np.trapz(flux_norm[condition], wavel[condition])/bin_width
    wavel_coarse = 0.5*(wavel_coarse[1:]+wavel_coarse[:-1])
    # Adjust the edges of half bin width
    half_width = [0.5*(wavel_coarse[1]-wavel_coarse[0]), 0.5*(wavel_coarse[-1]-wavel_coarse[-2])]
    wavel_coarse = np.hstack((wavel_coarse[0]-half_width[0], wavel_coarse, wavel_coarse[-1]+half_width[1]))
    flux_coarse = np.hstack((flux_coarse[0], flux_coarse, flux_coarse[-1]))
    return (wavel_coarse, flux_coarse)
    
def plot_absorption(
        base,   # frequency array or vel array or wavelength array
        flux,   # normalized flux
        figure = None, # matplotlib figure
        axis   = None, # matplotlib axis
        ):
    import matplotlib.pyplot as plt
    if figure is None:
        fig = plt.figure()
        ax = plt.gca()
    else:
        fig = figure
        ax = axis
    ax.step(base, flux, where="mid", color="tab:red", linewidth=3)
    
def generate_velocity_base(
        freq, # frequency array in Hz
        wav0, # rest wavelength of transition in angstrom
        ):
    c  = const.c.cgs.value
    freq0 = c/(wav0*1.0e-08)
    return ((freq0-freq)/freq0*c)/1.0e+05 # in km/s

def optical_depth_nu(
        freq,   # frequency array in Hz
        Temp,   # temperature in kelvin
        vlos,   # line of sight velocity in km/s
        column, # column density in cm^-2
        a_ion,  # atomic mass of the element in amu
        wav0,   # rest wavelength of transition in angstrom
        f_lu,   # oscillator stength in CGS
        gamma_ul, # sum of Einstein transition probability coefficients in s^-1
        ):
    # Generate optical depth at a higher resolution and deposit the average at lower resolution
    # This ensures same column width
    kB = const.k_B.cgs.value
    c  = const.c.cgs.value
    e_esu = const.e.esu.value
    amu = const.u.cgs.value
    me = const.m_e.cgs.value
    
    b = np.sqrt(2.*kB*Temp/(a_ion*amu))
    prefactor = (column*f_lu)*np.pi*e_esu**2/(me*c)

    freq0 = c/(wav0*1.0e-08)
    vlos_cms = vlos*1e+05
    freq0_doppler = freq0/np.sqrt((1.+vlos_cms/c)/(1.-vlos_cms/c))
    sigma = b/(np.sqrt(2)*c)
    gamma_nu = gamma_ul/(4.*np.pi*freq0_doppler)

    freq_norm = (freq0_doppler-freq)/freq0_doppler

    phi = voigt_profile(freq_norm,sigma,gamma_nu)/freq0_doppler # line profile
    optical_depth = prefactor*phi

    return ( freq, optical_depth )

def pepare_absorption(
        max_vel, # maximum velocity in km/s
        wav0,    # rest wavelength of transition in angstrom
        spectral_res, # spectral resolution in km/s
        ):
    c  = const.c.cgs.value/1.0e+05
    wav_range = wav0*(np.sqrt((1.+max_vel/c)/(1.-max_vel/c))-1.) # in angstrom
    del_wave  = wav0/(c/spectral_res)
    spec_wave = np.arange(wav0-wav_range, wav0+wav_range+del_wave, del_wave)
    return spec_wave  # in angstrom

def generate_absorption(
        max_vel, # maximum velocity in km/s
        spectral_res, # spectral resolution
        Temp,   # temperature in kelvin
        vlos,   # line of sight velocity in km/s
        column, # column density in cm^-2
        a_ion,  # atomic mass of the element in amu
        wav0,   # rest wavelength of transition in angstrom
        f_lu,   # oscillator stength in CGS
        gamma_ul, # sum of Einstein transition probability coefficients in s^-1
        ):
    c  = const.c.cgs.value
    spec_wave = pepare_absorption(max_vel, wav0, spectral_res) # in angstrom
    freq = c/(spec_wave*1.0e-08)
    return optical_depth_nu(freq, Temp, vlos, column, a_ion, wav0, f_lu, gamma_ul)