Generate synthetic absorption line spectra along a sightline

absorption_line_profile_calculator.py

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

@author: alankar

Absorption-line utilities for synthetic spectra along a sightline.

This module provides helper functions to build an absorption profile from
cell-wise gas properties (temperature, LOS velocity, ion density, and path
length), convert optical depth to normalized flux, optionally degrade to a
coarser spectral resolution, and compute equivalent widths (EW).

Typical workflow
----------------
1. For each LOS cell, call :func:`generate_absorption` and accumulate optical
    depth ``tau_eff``.
2. Convert accumulated optical depth to normalized flux with
    :func:`generate_norm_flux`.
3. Rebin the high-resolution flux to a coarse instrument grid using
    :func:`calc_flux_coarse`.
4. Plot high/low-resolution spectra in wavelength or velocity space
    (``plot_absorption`` + ``generate_velocity_base``).
5. Measure equivalent width with :func:`calculate_EW`.

Example
-------
The pattern below matches the intended usage for a single transition
(e.g. MgII 2796):

.. code-block:: python

     tau_eff = None
     column = 0.0
     for i in range(nH.shape[0]):
          temperature = Temp[i]
          del_x = dz[i]
          n_ion = n_MgII[i]
          vlos = velz[i] / 1.0e5  # km/s

          freq, tau_cell = generate_absorption(
                vel_range,
                spec_res[0],
                temperature,
                vlos,
                n_ion * del_x,
                a_ion,
                wave0,
                f_lu,
                gamma_ul,
          )
          tau_eff = tau_cell if tau_eff is None else (tau_eff + tau_cell)
          column += n_ion * del_x

     flux = generate_norm_flux(tau_eff)
     del_wave_coarse = wave0 / (const.c.cgs.value / 1.0e5 / spec_res[1])
     wave_angstrom = (const.c.cgs.value / freq) * 1.0e8
     wave_lowres, flux_lowres = calc_flux_coarse(wave_angstrom, flux, del_wave_coarse)
     ew_lowres = calculate_EW(wave_lowres, flux_lowres)
     ew_highres = calculate_EW(wave_angstrom, flux)

Notes
-----
- ``generate_absorption`` returns ``(freq, tau_nu)`` for one cell/transition.
- Summing ``tau_nu`` across cells is equivalent to multiplying transmission in
  intensity space.
- Units are CGS-based internally; function arguments document required units.
- Line data (rest wavelength, oscillator strength, Einstein A coefficient) must be
  provided for each transition of interest. 
  Database: 
    https://github.com/linetools/linetools/blob/master/linetools/data/lines/linelist.ascii
    https://physics.nist.gov/PhysRefData/ASD/lines_form.html
"""

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)