pythonradex 1.0.0__py3-none-any.whl

pythonradex/escape_probability_functions.py

@@ -0,0 +1,198 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Thu May 23 11:10:38 2024
+
+@author: gianni
+"""
+
+import numba as nb
+import numpy as np
+
+
+#It would be much more elegant to put these functions directly into the
+#classes defined in escape_probability.py. But I want to compile the beta functions
+#with numba, so I have to put these functions outside the class definitions
+#(defining them as staticmethod does not work here)
+
+min_reliable_tau = -1
+
+@nb.jit(nopython=True,cache=True)
+def clip_prob(prob):
+    return np.where(prob>1,1,np.where(prob<0,0,prob))
+
+####### functions for escape probability using Taylor expansion ######
+
+@nb.jit(nopython=True,cache=True)
+def identify_tau_regions(tau_nu):
+    tau_epsilon = 0.05 #important that this is not too small, to avoid the unstable region
+    normal_tau_region = tau_nu > tau_epsilon
+    small_tau_region = np.abs(tau_nu) <= tau_epsilon
+    negative_tau_region = (tau_nu >= min_reliable_tau) & (tau_nu<-tau_epsilon)
+    unreliable_negative_region = tau_nu < min_reliable_tau
+    total_region_points = normal_tau_region.sum()+small_tau_region.sum()\
+                          +negative_tau_region.sum()+unreliable_negative_region.sum()
+    assert total_region_points==len(tau_nu),\
+                          'selection needs to cover all points of tau_nu'
+    return normal_tau_region,small_tau_region,negative_tau_region,\
+                unreliable_negative_region
+
+@nb.jit(nopython=True,cache=True)
+def beta_analytical_uniform_sphere(tau_nu):
+    '''Computes the escape probability for a uniform sphere analytically,
+    given the optical depth tau_nu'''
+    #see the RADEX manual for this formula; derivation is found in the old
+    #Osterbrock (1974) book, appendix 2. Note that Osterbrock uses tau for
+    #radius, while I use it for diameter
+    return 1.5/tau_nu*(1-2/tau_nu**2+(2/tau_nu+2/tau_nu**2)*np.exp(-tau_nu))
+
+@nb.jit(nopython=True,cache=True)
+def beta_Taylor_uniform_sphere(tau_nu):
+    '''Computes the escape probability of a uniform sphere using a Taylor expansion,
+    given the optical depth tau_nu'''
+    #Taylor expansion of beta for uniform sphere is easier to evaluate numerically
+    #(for small tau_nu)
+    #Series calculated using Wolfram Alpha; not so easy analytically, to
+    #calculate the limit as tau->0, use rule of L'Hopital
+    return 1 - 0.375*tau_nu + 0.1*tau_nu**2 - 0.0208333*tau_nu**3
+
+@nb.jit(nopython=True,cache=True)
+def beta_analytical_LVG_slab(tau_nu):
+    #see line 357 in matrix.f and RADEX paper, or Scoville & Solomon (1974)
+    return (1-np.exp(-3*tau_nu))/(3*tau_nu)
+
+@nb.jit(nopython=True,cache=True)
+def beta_Taylor_LVG_slab(tau_nu):
+    return 1 - (3*tau_nu)/2 + (3*tau_nu**2)/2 - (9*tau_nu**3)/8
+
+#about the LVG sphere
+#we need to assume that the LVG sphere is homogeneous, because from formula 2.6.36
+#in Elitzur92, it becomes clear that the optical depth in a velocity-
+#coherent element depends on the local number density. But since pythonradex assumes
+#uniform excitation temperature and a single optical depth characerising the source,
+#the local number density must be the same everywhere
+#Thus, we need to assume that the LVG sphere is uniform, otherwise we cannot handle it
+#in pythonradex.
+#We are indeed considering the model by Goldreich & Kwan (1974) which has uniform density.
+#in fact, same argument also implies that LVG slab needs to be uniform...
+
+@nb.jit(nopython=True,cache=True)
+def beta_analytical_LVG_sphere(tau_nu):
+    #e.g Elitzur p. 44, or Ramos & Elitzur 2018, eq. 14
+    return (1-np.exp(-tau_nu))/tau_nu
+
+@nb.jit(nopython=True,cache=True)
+def beta_Taylor_LVG_sphere(tau_nu):
+    #from Wolfram Alpha
+    return 1 - tau_nu/2 + tau_nu**2/6 - tau_nu**3/24
+
+def generate_Taylor_beta(beta_ana,beta_Taylor):
+    @nb.jit(nopython=True,cache=True)
+    def beta(tau_nu):
+        #use Taylor expansion if tau is below epsilon; this is to avoid numerical
+        #problems with the anlytical formula to compute the escape probability.
+        #Concerning negative optical depth, the RADEX paper advises that
+        #negativ tau (inverted population) cannot be treated correctly by a
+        #non-local code like RADEX, and that results for lines with tau <~ -1
+        #should be ignored. Moreover, negative tau can make the code crash:
+        #very negative tau leads to very large transition rates, which makes the
+        #matrix of the rate equations ill-conditioned.
+        #thus, for tau < -1, I just take abs(tau). Note that this is different
+        #from RADEX: they make something quite strange: they have different
+        #approximations for positive large, normal and small tau. But then they
+        #use abs(tau) to decide which function to use, but then use tau in that
+        #function (line 333 in matrix.f)
+        normal,small,negative,unreliable = identify_tau_regions(tau_nu=tau_nu)
+        prob = np.empty_like(tau_nu)
+        prob[normal] = beta_ana(tau_nu[normal])
+        #here I use tau even if tau < 0:
+        prob[small] = beta_Taylor(tau_nu[small])
+        prob[negative] = beta_ana(tau_nu[negative])
+        #here I just use abs(tau) to stabilize the code
+        prob[unreliable] = beta_ana(np.abs(tau_nu[unreliable]))
+        assert np.all(np.isfinite(prob))
+        return clip_prob(prob)
+    return beta
+
+beta_uniform_sphere = generate_Taylor_beta(beta_ana=beta_analytical_uniform_sphere,
+                                           beta_Taylor=beta_Taylor_uniform_sphere)
+beta_LVG_slab = generate_Taylor_beta(beta_ana=beta_analytical_LVG_slab,
+                                     beta_Taylor=beta_Taylor_LVG_slab)
+beta_LVG_sphere = generate_Taylor_beta(beta_ana=beta_analytical_LVG_sphere,
+                                       beta_Taylor=beta_Taylor_LVG_sphere)
+
+######### functions for uniform slab  ##################
+#see Elitzur92 (https://ui.adsabs.harvard.edu/abs/1992ASSL..170.....E/abstract,
+#Problem 2.12) for confirmation of the formulas used here
+
+@nb.jit(nopython=True,cache=True)
+def integral_term_for_uniform_slab(tau):
+    mu = np.linspace(1e-5,1,200)
+    return np.trapezoid((1-np.exp(-tau/mu))*mu,mu)
+
+tau_grid_for_UniformSlab = np.logspace(-5,2,1000)
+min_grid_tau = np.min(tau_grid_for_UniformSlab)
+max_grid_tau = np.max(tau_grid_for_UniformSlab)
+#the expression for the flux contains an integral term;
+#here I pre-compute this term so it can be interpolated to speed up the code
+integral_term_grid = np.array([integral_term_for_uniform_slab(tau)
+                               for tau in tau_grid_for_UniformSlab])
+
+@nb.jit(nopython=True,cache=True)
+def interpolated_integral_term(tau):
+    interp = np.interp(x=tau,xp=tau_grid_for_UniformSlab,fp=integral_term_grid)
+    #limiting values:
+    # -very large tau: integral term goes to 0.5
+    # -very small tau: integral term goes to tau
+    #note that I cannot use the left and right keywords of np.interp because
+    #I want to compile this function; so I use np.where instead
+    return np.where(tau>max_grid_tau,0.5,np.where(tau<min_grid_tau,tau,interp))
+
+@nb.jit(nopython=True,cache=True)
+def beta_uniform_slab(tau_nu):
+    #for negative tau, this function will also return 1, which is fine if tau
+    #is close to 0, but not correct for very negative tau; in general, results
+    #should be ignored for tau < -1
+    int_term = interpolated_integral_term(tau=tau_nu)
+    prob = np.where(tau_nu<min_grid_tau,1,int_term/tau_nu)
+    assert np.all(np.isfinite(prob))
+    return clip_prob(prob)
+
+######### LVG sphere as implemented by RADEX ##################
+
+@nb.jit(nopython=True,cache=True)
+def beta_LVG_sphere_RADEX_gtr7(tau_nu):
+    tau_nu_r = tau_nu/2
+    return 2/(tau_nu_r*4*(np.sqrt(np.log(tau_nu_r/np.sqrt(np.pi)))))
+
+@nb.jit(nopython=True,cache=True)
+def beta_LVG_sphere_RADEX_less7(tau_nu):
+    tau_nu_r = tau_nu/2
+    return 2*(1-np.exp(-2.34*tau_nu_r))/(4.68*tau_nu_r)
+
+@nb.jit(nopython=True,cache=True)
+def beta_LVG_sphere_RADEX(tau_nu):
+    #RADEX paper, eq. 18, claims that they use beta_analytical_LVG_sphere
+    #(beta= (1-exp(-tau))/tau
+    #but in fact they use a different formula from de Jong+80
+    #see line 346 in matrix.f
+    #I use the same formulae as RADEX, except that I take care of negative tau
+    #handling negative tau turns out to be important, e.g. for CO, Tkin=200,
+    #width_v=3 km/s, Ntot=1e20 cm-2, RADEX gives an invalid solution
+    assert -7 < min_reliable_tau < -0.01
+    gtr7 = 7 <= tau_nu
+    less7 = (0.001 <= tau_nu) & (tau_nu < 7)
+    small = (-0.001 <= tau_nu) & (tau_nu < 0.001)
+    negative = (min_reliable_tau <= tau_nu) & (tau_nu < -0.001)
+    unreliable_less7 = (-7 <= tau_nu) & (tau_nu < min_reliable_tau)
+    unreliable_gtr7 = tau_nu < -7
+    assert gtr7.sum()+less7.sum()+small.sum()+negative.sum()+unreliable_less7.sum()\
+              +unreliable_gtr7.sum() == tau_nu.size
+    beta = np.empty_like(tau_nu)
+    beta[gtr7] = beta_LVG_sphere_RADEX_gtr7(tau_nu[gtr7])
+    beta[less7] = beta_LVG_sphere_RADEX_less7(tau_nu[less7])
+    beta[small] = 1.0
+    beta[negative] = beta_LVG_sphere_RADEX_less7(tau_nu[negative])
+    beta[unreliable_less7] = beta_LVG_sphere_RADEX_less7(np.abs(tau_nu[unreliable_less7]))
+    beta[unreliable_gtr7] = beta_LVG_sphere_RADEX_gtr7(np.abs(tau_nu[unreliable_gtr7]))
+    return clip_prob(beta)

pythonradex/flux.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Aug  2 10:32:03 2024
+
+@author: gianni
+"""
+
+from scipy import constants
+from pythonradex import helpers
+import numpy as np
+import warnings
+
+
+class FluxCalculator():
+
+    tau_peak_fraction = 1e-2
+    nu_per_FHWM = 20
+    n_nu_elements = {'regular':15,'LVG sphere':51}
+
+    def __init__(self,emitting_molecule,level_population,geometry_name,
+                 V_LVG_sphere,compute_flux_nu,tau_nu0_individual_transitions,
+                 tau_dust,S_dust):
+        self.emitting_molecule = emitting_molecule
+        self.level_population = level_population
+        self.geometry_name = geometry_name
+        self.V_LVG_sphere = V_LVG_sphere
+        self.Tex = self.emitting_molecule.get_Tex(level_population=level_population)
+        self.compute_flux_nu = compute_flux_nu
+        self.tau_nu0_individual_transitions = tau_nu0_individual_transitions
+        self.tau_dust = tau_dust
+        self.S_dust = S_dust
+        self.is_LVG_sphere = self.geometry_name == 'LVG sphere'
+
+    def get_flux_parameters_for_rectangular_flux(self,transitions):
+        n_nu_elements = self.n_nu_elements['LVG sphere'] if self.is_LVG_sphere\
+                        else self.n_nu_elements['regular']
+        assert n_nu_elements%2 == 1,\
+                         'need odd number of elements to include nu0 in the array'
+        middle_index = n_nu_elements//2
+        nu0 = self.emitting_molecule.nu0[transitions]
+        width_nu = self.emitting_molecule.width_v/constants.c*nu0
+        # a matrix where first index is transition, second index is nu
+        nu = np.linspace(nu0-width_nu/2,nu0+width_nu/2,n_nu_elements,axis=-1)
+        tau_nu = 1/nu**2
+        tau_nu0 = self.tau_nu0_individual_transitions[transitions]
+        tau_nu *= tau_nu0[:,None]/tau_nu[:,middle_index][:,None]
+        source_function = helpers.B_nu(T=self.Tex[transitions][:,None],nu=nu)
+        return nu,tau_nu,source_function
+
+    def fast_line_fluxes_rectangular_without_overlap(self,solid_angle,transitions):
+        #note that integrated fluxes only make sense for non-overlapping lines,
+        #so this function can only be used for non-overlapping lines
+        nu,tau_nu,source_function = self.get_flux_parameters_for_rectangular_flux(
+                                          transitions=transitions)
+        kwargs = {'tau_nu':tau_nu,'source_function':source_function,
+                  'solid_angle':solid_angle}
+        if self.is_LVG_sphere:
+            kwargs['nu'] = nu
+            kwargs['nu0'] = self.emitting_molecule.nu0[transitions][:,None]
+            kwargs['V'] = self.V_LVG_sphere
+        flux_nu = self.compute_flux_nu(**kwargs)
+        flux = np.trapezoid(flux_nu,nu,axis=1)
+        return flux
+
+    def fast_line_fluxes_Gaussian_without_overlap(self,solid_angle,transitions):
+        nu0 = self.emitting_molecule.nu0[transitions]
+        FWHM_nu = self.emitting_molecule.width_v/constants.c*nu0
+        sigma_nu = helpers.FWHM2sigma(FWHM_nu)
+        #for thin emission, I just want to include out to a certain fraction
+        #of the peak
+        #but for thick emission, the spectrum is saturated, so a fraction of the
+        #peak is not useful; in those cases, I need to set an absolute value
+        #for the minimum tau to include
+        tau_nu0 = self.tau_nu0_individual_transitions[transitions]
+        peak_frac = self.tau_peak_fraction*tau_nu0
+        min_tau = np.where(peak_frac>0.01,0.01,peak_frac)
+        #calculate the distance in freq spac between tau_nu0 and min_tau:
+        #(min_tau = tau_nu0*exp(-(nu-nu0)**2/(2*sigma_nu**2))
+        Delta_nu = sigma_nu*np.sqrt(-2*np.log(min_tau/tau_nu0))
+        n_nu = np.max((2*Delta_nu/FWHM_nu * self.nu_per_FHWM).astype(int))
+        #first index transition, second index frequency:
+        nu = np.linspace(nu0-Delta_nu,nu0+Delta_nu,n_nu,axis=-1)
+        #doing an approximation: in principle, tau has an additional 1/nu**2 dependence,
+        #but if Delta_nu is small compared to nu0, that dependence is negligible
+        phi_nu_shape = np.exp(-(nu-nu0[:,None])**2/(2*sigma_nu[:,None]**2))
+        tau_nu = tau_nu0[:,None]*phi_nu_shape
+        source_function = helpers.B_nu(T=self.Tex[:,None][transitions],nu=nu)
+        spectrum = self.compute_flux_nu(tau_nu=tau_nu,source_function=source_function,
+                                        solid_angle=solid_angle)
+        flux = np.trapezoid(spectrum,nu,axis=1)
+        return flux
+
+    def all_transitions(self):
+        return list(range(self.emitting_molecule.n_rad_transitions))
+
+    def determine_tau_of_overlapping_lines(self,transitions):
+        total_tau = []
+        for i in transitions:
+            tau = 0
+            for j in self.emitting_molecule.overlapping_lines[i]:
+                tau += self.tau_nu0_individual_transitions[j]
+            total_tau.append(tau)
+        return np.array(total_tau)
+
+    def fluxes_of_individual_transitions(self,solid_angle,transitions=None):
+        max_acceptable_tau = 0.1 #for dust and overlapping lines
+        if transitions is None:
+            transitions = self.all_transitions()
+        if self.emitting_molecule.any_line_has_overlap(transitions):
+            tau_overlapping_lines = self.determine_tau_of_overlapping_lines(
+                                                        transitions=transitions)
+            if np.any(tau_overlapping_lines>max_acceptable_tau):
+                raise ValueError('fluxes of individual lines can only be calculated'
+                                 +' for non-overlapping lines or thin overlapping lines')
+        if np.any(self.tau_dust(self.emitting_molecule.nu0) > max_acceptable_tau):
+            raise ValueError('dust is not optically thin, cannot calculate'+
+                             ' fluxes of individual lines')
+        if self.emitting_molecule.line_profile_type == 'Gaussian':
+            fast_flux = self.fast_line_fluxes_Gaussian_without_overlap(
+                                solid_angle=solid_angle,transitions=transitions)
+        elif self.emitting_molecule.line_profile_type == 'rectangular':
+            fast_flux = self.fast_line_fluxes_rectangular_without_overlap(
+                                        solid_angle=solid_angle,transitions=transitions)
+        else:
+            raise ValueError(f'line profile {self.emitting_molecule.line_profile_type} '
+                             +'unknown')
+        return np.squeeze(fast_flux)
+
+    def set_nu(self,nu):
+        self.nu = nu
+        self.nu_selected_lines,self.nu_selected_line_indices = self.identify_lines()
+        self.tau_dust_nu = self.tau_dust(self.nu)
+        self.set_tau_nu_lines()
+        self.set_tau_nu_tot()
+
+    def identify_lines(self):
+        nu_min,nu_max = np.min(self.nu),np.max(self.nu)
+        selected_lines = []
+        selected_line_indices = []
+        for i,line in enumerate(self.emitting_molecule.rad_transitions):
+            if nu_min <= line.nu0 and line.nu0 <= nu_max:
+                selected_line_indices.append(i)
+                selected_lines.append(line)
+        return selected_lines,selected_line_indices
+
+    def construct_tau_nu_individual_line(self,line,nu,tau_nu0):
+        #tau is prop to phi_nu/nu**2
+        normalisation = line.line_profile.phi_nu0/line.nu0**2 / tau_nu0
+        return line.line_profile.phi_nu(nu)/nu**2 / normalisation
+
+    def set_tau_nu_lines(self):
+        self.tau_nu_lines = []
+        for line_index,line in zip(self.nu_selected_line_indices,self.nu_selected_lines):
+            tau = self.construct_tau_nu_individual_line(
+                              line=line,nu=self.nu,
+                              tau_nu0=self.tau_nu0_individual_transitions[line_index])
+            self.tau_nu_lines.append(tau)
+
+    def set_tau_nu_tot(self):
+        #note that this function automatically takes overlapping lines into account
+        #note also that I cannot use the get_tau_tot_nu method of the
+        #emitting_molecule, because that method only considers a single line
+        #and its overlaps
+        self.tau_nu_tot = np.sum(self.tau_nu_lines,axis=0) + self.tau_dust_nu
+
+    def get_S_tot(self):
+        S_nu = np.zeros_like(self.nu)
+        for i,line in enumerate(self.nu_selected_lines):
+            x1 = self.level_population[line.low.number]
+            x2 = self.level_population[line.up.number]
+            S_line = line.source_function(x1=x1,x2=x2)
+            S_nu += self.tau_nu_lines[i]*S_line
+        S_nu += self.S_dust(self.nu)*self.tau_dust_nu
+        S_tot = np.where(self.tau_nu_tot==0,0,S_nu/self.tau_nu_tot)
+        return S_tot
+
+    def spectrum(self,solid_angle):
+        #note that this function gives the right answer also in cases where
+        #line overlap treatment is not necessary
+        if self.is_LVG_sphere:
+            #assumption for LVG sphere: lines do not overlap and there is no dust,
+            #so I can just sum them up
+            #this also works for overlapping lines in the optically thin regime
+            #reason for this special case: the formula to compute the flux for an LVG
+            #sphere depends on nu0, which is not well defined if there are several
+            #lines
+            if self.emitting_molecule.any_line_has_overlap(
+                                line_indices=self.nu_selected_line_indices):
+                warnings.warn('LVG sphere geometry: lines are overlapping, '
+                              +'output spectrum will only be correct if all lines'
+                              +' are optically thin!')
+            for dust_func in (self.S_dust,self.tau_dust):
+                assert np.all(dust_func(self.nu)==0),'LVG does not support dust'
+            spectrum = np.zeros_like(self.nu)
+            for line_index,line,tau_nu_line in zip(self.nu_selected_line_indices,
+                                                   self.nu_selected_lines,
+                                                   self.tau_nu_lines):
+                LVG_sphere_kwargs = {'nu':self.nu,'nu0':line.nu0,'V':self.V_LVG_sphere}
+                source_function = helpers.B_nu(T=self.Tex[line_index],nu=self.nu)
+                spectrum += self.compute_flux_nu(
+                                    tau_nu=tau_nu_line,source_function=source_function,
+                                    solid_angle=solid_angle,**LVG_sphere_kwargs)
+        else:
+            source_function = self.get_S_tot()
+            spectrum = self.compute_flux_nu(
+                               tau_nu=self.tau_nu_tot,source_function=source_function,
+                               solid_angle=solid_angle)
+        return spectrum

pythonradex/helpers.py

@@ -0,0 +1,69 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+from scipy import constants
+import numba as nb
+
+
+@nb.jit(nopython=True,cache=True)
+def B21(A21,nu):
+    return constants.c**2/(2*constants.h*nu**3)*A21
+
+@nb.jit(nopython=True,cache=True)
+def B12(A21,nu,g1,g2):
+    return g2/g1*B21(A21=A21,nu=nu)
+
+@nb.jit(nopython=True,cache=True) 
+def B_nu(nu,T):
+    r"""Planck function (black body)
+    
+    Args:
+        nu (:obj:`float` or numpy.ndarray): frequency in [Hz]
+        T (:obj:`float` or numpy.ndarray): temperature in [K]
+
+    Returns:
+        float or numpy.ndarray: value of Planck function in [W/m\ :sup:`2`/Hz/sr]
+    """
+    return (2*constants.h*nu**3/constants.c**2
+           *(np.exp(constants.h*nu/(constants.k*T))-1)**-1)
+
+def generate_CMB_background(z=0):
+    r'''generates a function that gives the CMB background at redshift z
+    
+    Args:
+        z: (:obj:`float`): redshift
+
+    Returns:
+        func: a function giving CMB background in [W/m\ :sup:`2`/Hz/sr] for an input frequency in [Hz]
+        '''
+    T_CMB = 2.73*(1+z)
+    def CMB_background(nu):
+        return B_nu(nu=nu,T=T_CMB)
+    return CMB_background
+
+@nb.jit(nopython=True,cache=True) 
+def FWHM2sigma(FWHM):
+    """Convert FWHM of a Gaussian to standard deviation.
+    
+    Args:
+        FWHM (:obj:`float` or numpy.ndarray): FWHM of the Gaussian
+    
+    Returns:
+        float or numpy.ndarray: the standard deviation corresponding to the input FWHM"""
+    return FWHM/(2*np.sqrt(2*np.log(2)))
+
+@nb.jit(nopython=True,cache=True) 
+def relative_difference(a,b):
+    """Computes the elementwise relative difference between a and b.
+    In general, return |a-b|/a.
+    Special cases:
+    a=0 and b=0: return 0
+    a=0 and b!=0: return 1"""
+    abs_diff = np.abs(a-b)
+    rel_diff = np.where((a==0) & (b==0),0,np.where(a==0,1,abs_diff/a))
+    assert not np.any(np.isnan(rel_diff))
+    return np.abs(rel_diff)
+
+@nb.jit(nopython=True,cache=True)
+def Delta_nu(Delta_v,nu0):
+    return nu0 * Delta_v / constants.c