import os
import sys
import inspect
import numpy as np
from dataclasses import dataclass, fields, is_dataclass
from typing import Optional, Literal, Any
from astromugs.utils.struct import DiskParams
from astromugs.utils import custom
from astromugs.constants.constants import (
mu,
autocm,
amu,
Ggram,
kb,
M_sun,
)
#___________________________________________
# _______ __ _______ __ ___
# | __ \ | | | _____| | | / /
# | | | | | | | |____ | |/ /
# | | | | | | |____ | | |
# | |__| | | | ____| | | |\ \
# |_______/ |__| |_______| |__| \__\
#___________________________________________
[docs]
class Disk:
"""Static flared disk model for radiative transfer and chemistry simulations.
Provides methods to compute gas and dust physical quantities (density,
temperature, scale height, surface density, extinction) on various
coordinate grids.
Parameters
----------
params : DiskParams
Dataclass holding all disk model parameters.
dust : object, optional
Dust model object providing grain size distributions and opacities.
"""
[docs]
def __init__(
self,
params: DiskParams,
dust: Optional[object] = None
):
"""Initialize the Disk model.
Parameters
----------
params : DiskParams
Dataclass holding all disk model parameters. Its fields are also
copied as instance attributes for convenience.
dust : object, optional
Dust model object providing grain size distributions and opacities.
"""
self.params = params
self.dust = dust
# copy dataclass attributes to instance variables (optional)
for field in params.__dataclass_fields__:
setattr(self, field, getattr(params, field))
# Pre-load custom temperature table if available
self._custom_r = None
self._custom_Tgas = None
if params.sigma_compute == 'custom' and params.sigma_path is not None:
r_custom, _, _, T_gas = custom.surfacedensities(params.sigma_path)
if T_gas is not None:
self._custom_r = r_custom
self._custom_Tgas = T_gas
[docs]
def scaleheight(self, r):
"""Compute the gas scale height.
Parameters
----------
r : float or array_like
Distance(s) from the star in AU.
Returns
-------
hgas : float or ndarray
Gas scale height in cm.
Notes
-----
When a custom temperature profile is available (from the custom
surface density file), the scale height is computed directly from
``H = sqrt(k_B T r^3 / (mu m_H G M_*))`` at each radius.
Otherwise, a power-law profile is used via ``h0`` or ``tmidplan_ref``.
"""
r_arr = np.atleast_1d(r)
if self._custom_Tgas is not None:
# Interpolate T(r) from the custom table, then compute H(r) directly
T = self.temp_mid(r_arr)
hgas = np.sqrt((kb * T * (r_arr * autocm)**3) /
(mu * amu * Ggram * self.params.star_mass * M_sun))
else:
h_exp = (3./2.) - (self.params.q_exp/2.)
if self.params.h0 != None:
hgas = self.params.h0*autocm*(r_arr/(self.params.ref_radius))**h_exp
elif self.params.tmidplan_ref != None:
h0 = np.sqrt((kb*self.params.tmidplan_ref*(self.params.ref_radius*autocm)**3)/(mu*amu*Ggram*self.params.star_mass*M_sun))
self.params.h0 = h0/autocm
hgas = h0*(r_arr/(self.params.ref_radius))**h_exp
# Return scalar if input was scalar
if np.ndim(r) == 0:
return hgas.item() if hgas.size == 1 else hgas
return hgas
[docs]
def surfacedensity(self, r):
"""Compute the gas surface density.
Parameters
----------
r : float or array_like
Distance(s) from the star in AU.
Returns
-------
sigma_g : float or ndarray
Gas surface density in g/cm^2.
"""
if self.params.sigma_compute == 'custom':
r_custom, _, siggas_table, _ = custom.surfacedensities(self.params.sigma_path)
r_flat = np.clip(r.ravel() if hasattr(r, 'ravel') else np.atleast_1d(r), r_custom[0], r_custom[-1])
sigma_g = np.interp(r_flat, r_custom, siggas_table)
if hasattr(r, 'shape'):
sigma_g = sigma_g.reshape(r.shape)
return sigma_g
elif self.params.sigma_gas_ref != None:
sigma_g = 2*self.params.sigma_gas_ref*(r/(self.params.ref_radius))**(-self.params.p_exp) #2 or not??
else:
sigma_g = (1/self.params.dtogas)*self.sigma_d0*(r/(self.params.ref_radius))**(-self.params.p_exp)
return sigma_g
[docs]
def omega2(self, r):
"""Compute the square of the Keplerian angular velocity.
Parameters
----------
r : float or array_like
Distance(s) from the star in AU.
Returns
-------
float or ndarray
Keplerian angular velocity squared in s^-2.
"""
return (Ggram*self.params.star_mass*M_sun)/(r*autocm)**3
[docs]
def temp_mid(self, r):
"""Compute the midplane temperature radial profile.
Parameters
----------
r : float or array_like
Distance(s) from the star in AU.
Returns
-------
float or ndarray
Midplane temperature in K.
Notes
-----
When a custom temperature profile is loaded from the surface density
file, the temperature is interpolated from that table. Otherwise, the
analytic power law ``T = T_ref * (r/r_ref)^(-q)`` is used.
"""
if self._custom_Tgas is not None:
r_flat = np.atleast_1d(r).ravel()
r_clipped = np.clip(r_flat, self._custom_r[0], self._custom_r[-1])
T = np.interp(r_clipped, self._custom_r, self._custom_Tgas)
if np.ndim(r) == 0:
return T.item()
return T.reshape(np.shape(r))
return self.params.tmidplan_ref*(r/self.params.ref_radius)**(-self.params.q_exp)
[docs]
def temp_atmos(self, r):
"""Compute the atmospheric temperature radial profile.
Parameters
----------
r : float or array_like
Distance(s) from the star in AU.
Returns
-------
float or ndarray
Atmospheric temperature in K.
"""
return self.params.tatmos_ref*(r/self.params.ref_radius)**(-self.params.q_exp)
[docs]
def temp_altitude(self, r, z):
"""Compute the vertical temperature profile following Williams & Best (2014).
Parameters
----------
r : array_like
Radial distances from the star in AU.
z : array_like
Normalized altitude coordinates (in units of scale height).
Returns
-------
ndarray
Temperature array of shape ``(len(r), len(z))`` in K.
"""
hg = self.scaleheight(r)
tmid = self.temp_mid(r)
tatm = self.temp_atmos(r)
hhg, zz = np.meshgrid(hg, z, indexing='ij')
ttmid, zz = np.meshgrid(tmid, z, indexing='ij')
ttatm, zz = np.meshgrid(tatm, z, indexing='ij')
zz = hhg*zz
zz0, z = np.meshgrid(zz[:, 0], z, indexing='ij')
return ttmid+(ttatm-ttmid)*np.sin((np.pi*zz)/(2*zz0))**(2*self.params.sigma_t)
[docs]
def verticaldensity_gauss(self, r):
"""Compute the gas number density assuming a Gaussian vertical profile.
Parameters
----------
r : array_like
Radial distances from the star in AU.
Returns
-------
ndarray
Gas number density in cm^-3.
Notes
-----
Valid only for vertically isothermal disks. The surface density is
divided by the mean molecular mass to yield number density rather
than mass density.
"""
sigma = self.surfacedensity()
sigma[(r >= self.params.rout) ^ (r <= self.params.rin)] = 0e0
H = self.scaleheight(r)
z = self.altitudes()
return (sigma/(mu*amu*H*autocm*np.sqrt(2.*np.pi)))\
*np.exp(-(z)**2./(2.*(H)**2.))
[docs]
def density(self, x1, x2, x3=None):
"""Compute the gas mass density assuming hydrostatic equilibrium.
Parameters
----------
x1 : array_like
Radial coordinate in AU (spherical: radii; nautilus: radii).
x2 : array_like
Second coordinate (spherical: colatitude in rad; nautilus:
normalized altitudes in units of scale height).
x3 : array_like, optional
Third coordinate (spherical: azimuthal angle in rad). Required
for spherical grids.
Returns
-------
rhog : ndarray
Gas mass density in g/cm^3.
Notes
-----
The Gaussian vertical profile is an isothermal approximation. For a
non-vertically isothermal profile, the density can be computed
iteratively and slightly deviates from a Gaussian profile.
"""
if self.params.coordsystem == "spherical":
rhog = np.ones((len(x1), len(x2), len(x3)))
rt, tt, pp = np.meshgrid(x1*autocm, x2, x3, indexing='ij')
rr = rt*np.sin(tt)
zz = rt*np.cos(tt)
zzmax = self.params.rout*autocm*np.cos(tt)
border = np.greater_equal(abs(zzmax), abs(zz))*1 #used to spherically cut the disk in case the grid is larger than the disk outer radius e.g. if an envelope is present
sigma_g = self.surfacedensity(rr/autocm)
hg = self.scaleheight(rr/autocm)
rhog = (sigma_g/(np.sqrt(2*np.pi)*hg))*np.exp(-(zz[:,:,:]**2)/(2*hg**2))*border
if self.params.coordsystem == "nautilus":
omega2 = self.omega2(x1)
sigma_g = self.surfacedensity(x1)
hg = self.scaleheight(x1)
temp = self.temp_altitude(x1, x2)
midplane_dens = (sigma_g)/(hg*np.sqrt(2*np.pi))
rhog = np.ones((len(x1), len(x2)))
rhog[:, 0] = 0 # init
for r in range(0, len(x1), 1):
for z in range(1, len(x2), 1):
rhog[r, z] = rhog[r, z-1] - (np.log(temp[r,z]) - np.log(temp[r,z-1])) - ((mu*amu*omega2[r]*hg[r]*x2[z]*(hg[r]*x2[z] - hg[r]*x2[z-1]))/(kb*temp[r,z]))
rhog[r] = np.exp(rhog[r])
maximum=np.amax(rhog[r])
rhog[r] = (rhog[r]*midplane_dens[r])/maximum
return rhog
[docs]
def numberdensity(self, x1, x2, x3=None):
"""Compute the gas number density assuming hydrostatic equilibrium.
Parameters
----------
x1 : array_like
Radial coordinate in AU.
x2 : array_like
Second coordinate (colatitude in rad or normalized altitude).
x3 : array_like, optional
Azimuthal angle in rad (spherical grids only).
Returns
-------
ng : ndarray
Gas number density in cm^-3.
Notes
-----
If vertically isothermal, the profile is Gaussian. Otherwise the
density is computed iteratively and deviates from a Gaussian.
"""
ng = self.density(x1, x2, x3)/(mu*amu)
return ng
[docs]
def viscous_accretion_heating(self, x1, x2, x3=None):
"""Compute the viscous accretion heating rate.
Parameters
----------
x1 : array_like
Radial coordinate in AU.
x2 : array_like
Colatitude in rad.
x3 : array_like, optional
Azimuthal angle in rad.
Returns
-------
q_visc : ndarray
Volumetric viscous heating rate in erg/(cm^3 s).
Notes
-----
Currently implemented for spherical coordinates only. The heating is
distributed vertically as a Gaussian and truncated beyond
``lim_h`` scale heights.
"""
if self.params.coordsystem == "spherical":
yrtosec = 31536000 # number of second in a year.
q_visc = np.ones((len(x1), len(x2), len(x3)))
rt, tt, pp = np.meshgrid(x1*autocm, x2, x3, indexing='ij')
rr = rt*np.sin(tt)
zz = rt*np.cos(tt)
zzmax = self.params.rout*autocm*np.cos(tt)
sigma_g = self.surfacedensity(rr/autocm)
hg = self.scaleheight(rr/autocm)
omega2 = self.omega2(rr/autocm)
border_max = np.greater_equal(abs(zzmax), abs(zz))*1 # make border at rout
border_h = np.greater_equal(abs(self.params.lim_h*hg), abs(zz))*1 #make border at lim_h scale height.
border = border_max*border_h
q_visc = ((3*self.params.acc_rate*M_sun*omega2)/(4*np.pi*np.sqrt(2*np.pi)*hg*yrtosec))*np.exp(-(zz[:,:,:]**2)/(2*hg**2))*border
#q_visc[q_visc==0.0] = 1e-30
return q_visc
[docs]
def surfacedensity_d(self, r):
"""Compute the dust surface density per grain-size bin.
Parameters
----------
r : float or array_like
Distance(s) from the star in AU.
Returns
-------
sigmad : ndarray
Dust surface density for each grain-size bin in g/cm^2, shape
``(n_sizes, *r.shape)``.
sig_single : float or ndarray
Total (summed over sizes) dust surface density in g/cm^2.
"""
if self.params.sigma_compute == 'model':
sigmad = []
fraction = self.dust.massfraction()
#print(r.shape)
#sig = np.ones((fraction.shape, r[0].shape, r[1].shape, r[2].shape))
#print(sig)
if self.params.sigma_gas_ref != None:
sigma_di0 = fraction*self.params.dtogas*self.params.sigma_gas_ref
sigma_single0 = self.params.dtogas*self.params.sigma_gas_ref
else:
self.sigma_d0 = (2-self.params.p_exp)*self.params.dust_mass*M_sun/(2*np.pi*(self.params.ref_radius)**(self.params.p_exp)) / \
(self.params.rout**(-self.params.p_exp+2) - self.params.rin**(-self.params.p_exp+2)) /autocm**2
sigma_di0 = fraction*self.sigma_d0
sigma_single0 = self.sigma_d0
for s in sigma_di0:
sig = s*(r/self.params.ref_radius)**(-self.params.p_exp)
sig[(r >= self.params.rout) ^ (r <= self.params.rin)] = 0e0 # In case sigma is outside disk outer boundaries
sigmad.append(sig)
sig_single = sigma_single0*(r/self.params.ref_radius)**(-self.params.p_exp)
sig_single[(r >= self.params.rout) ^ (r <= self.params.rin)] = 0e0 # In case sigma is outside disk outer boundaries
return np.array(sigmad), sig_single
elif self.params.sigma_compute == 'custom':
r_custom, sigmad_table, siggas_table, _ = custom.surfacedensities(self.params.sigma_path)
# sigmad_table: (n_sizes, n_file_radii), r_custom: (n_file_radii,)
# Interpolate each size bin onto the grid radii r (can be 2D or 3D meshgrid)
# Clamp radii to the table range to avoid midplane band artifacts
# in spherical grids (the outer cutoff is handled by border in density_d)
r_flat = np.clip(r.ravel(), r_custom[0], r_custom[-1])
sigmad = np.array([
np.interp(r_flat, r_custom, sigmad_table[i]).reshape(r.shape)
for i in range(sigmad_table.shape[0])
])
sig_single = np.interp(r_flat, r_custom, sigmad_table.sum(axis=0)).reshape(r.shape)
return sigmad, sig_single
[docs]
def scaleheight_d(self, r):
"""Compute the dust scale height for each grain size.
Parameters
----------
r : array_like
Radial distances from the star in AU.
Returns
-------
hd : ndarray
Dust scale height for each grain size in cm, shape
``(n_sizes, *r.shape)``.
hd_single : float or ndarray
Dust scale height for the single representative grain in cm.
Notes
-----
When settling is enabled, the dust scale height is reduced relative
to the gas scale height based on the Stokes number, the Schmidt
number, and the turbulent alpha parameter.
"""
hd = []
hg = self.scaleheight(r)
sizes = self.dust.sizes() #grain size in microns
rsingle = self.dust.rsingle
if self.params.settling == True:
sigma_g = self.surfacedensity(r)
for a in sizes[-1]:
stoptime_mid =(np.pi*a*1e-4*self.params.rho_m)/(2*sigma_g)
hd.append(hg/(np.sqrt(1 + stoptime_mid*(self.params.schmidtnumber/self.params.alpha))))
stoptime_mid_single =(np.pi*self.dust.rsingle*1e-4*self.params.rho_m)/(2*sigma_g)
hd_single = hg/(np.sqrt(1 + stoptime_mid_single*(self.params.schmidtnumber/self.params.alpha)))
return np.array(hd), hd_single
if self.params.settling == False:
for a in sizes[-1]:
hd.append(hg)
return np.array(hd), hg
[docs]
def density_d(self, x1, x2, x3=None):
"""Compute the dust mass density for each grain-size bin.
Parameters
----------
x1 : array_like
Radial coordinate in AU.
x2 : array_like
Second coordinate (colatitude in rad or normalized altitude).
x3 : array_like, optional
Azimuthal angle in rad (spherical grids only).
Returns
-------
rhod : ndarray
Dust mass density in g/cm^3, shape
``(n_sizes, len(x1), len(x2)[, len(x3)])``.
"""
rhod = []
rhod_single = []
if self.params.coordsystem =='spherical':
rt, tt, pp = np.meshgrid(x1*autocm, x2, x3, indexing='ij')
rr = rt*np.sin(tt)
zz = rt*np.cos(tt)
zzmax = self.params.rout*autocm*np.cos(tt)
sigmad, sigmad_single = self.surfacedensity_d(rr/autocm)
hd, hd_single = self.scaleheight_d(rr/autocm)
rhod = np.ones( (len(hd), len(x1), len(x2), len(x3)))
rhod_single = np.ones((len(x1), len(x2), len(x3)))
border = np.greater_equal(abs(zzmax), abs(zz))*1 #used to spherically cut the disk in case the grid is larger than the disk outer radius e.g. if an envelope is present. *1 to convert boolean to 0,1 values.
for i in range(len(hd)): #len(hd) is equal to the number of grain sizes
rhod[i, :, :, :] = (sigmad[i,:,:,:]/(np.sqrt(2*np.pi)*hd[i,:,:,:]))*np.exp(-(zz[:,:,:]**2)/(2*hd[i,:,:,:]**2))*border
if self.params.coordsystem =='nautilus':
rr, zz = np.meshgrid(x1*autocm, x2, indexing='ij')
sigmad, sigmad_single = self.surfacedensity_d(rr/autocm)
hg = self.scaleheight(rr/autocm)
hd, hd_single = self.scaleheight_d(rr/autocm)
rhod = np.ones((len(hd), len(x1), len(x2)))
rhod_single = np.ones((len(x1), len(x2)))
zz = hg*zz
for i in range(len(hd)): #len(hd) is equal to the number of grain sizes
rhod[i, :, :] = (sigmad[i]/(np.sqrt(2*np.pi)*hd[i]))*np.exp(-((zz)**2)/(2*hd[i]**2))
rhod_single[:, :] = (sigmad_single/(np.sqrt(2*np.pi)*hd_single))*np.exp(-((zz)**2)/(2*hd_single**2))
self.rhod_single = rhod_single
return rhod
[docs]
def numberdensity_d(self, x1, x2, x3=None):
"""Compute the dust number density for each grain-size bin.
Parameters
----------
x1 : array_like
Radial coordinate in AU.
x2 : array_like
Second coordinate (colatitude in rad or normalized altitude).
x3 : array_like, optional
Azimuthal angle in rad (spherical grids only).
Returns
-------
dens : ndarray
Dust number density in cm^-3, shape
``(n_sizes, len(x1), len(x2)[, len(x3)])``.
Examples
--------
Access densities at the third radius for the second grain species:
>>> n_d = disk.numberdensity_d(x1, x2, x3)
>>> n_d[1, 2, :]
"""
mass = self.dust.grainmass()
dens = self.density_d(x1, x2, x3)
for i in range(len(mass)):
dens[i] = dens[i]/mass[i]
return dens
[docs]
def numberdensity_d_single(self, x1, x2, x3=None):
"""Compute the dust number density for a single representative grain.
Useful when RADMC-3D uses multiple grain sizes but the chemistry
solver (e.g. Nautilus) requires a single grain population.
Parameters
----------
x1 : array_like
Radial coordinate in AU.
x2 : array_like
Second coordinate (colatitude in rad or normalized altitude).
x3 : array_like, optional
Azimuthal angle in rad (spherical grids only).
Returns
-------
dens_single : ndarray
Dust number density in cm^-3, shape ``(len(x1), len(x2))``.
Notes
-----
If ``nb_species == 1``, the result equals ``numberdensity_d[0]``.
"""
mass_single = self.dust.grainmass_single()
dens_single = self.rhod_single
dens_single = dens_single/mass_single
return dens_single
[docs]
def q_ext(self, lam, A=2, q_c=4):
"""Compute the extinction efficiency for each grain size and wavelength.
Parameters
----------
lam : array_like
Wavelengths in microns.
A : float, optional
Extinction efficiency in the geometric-optics regime (default 2).
q_c : float, optional
Extinction efficiency at the critical wavelength (default 4).
Returns
-------
qext : ndarray
Extinction efficiency, shape ``(n_sizes, len(lam))``.
Dimensionless.
"""
sizes = self.dust.sizes()
lambda_c = 2*np.pi*sizes[-1]
qext = np.ones(( len(sizes[-1]), len(lam) ))
for idx_a, a in enumerate(sizes[-1]):
for idx_wl, wl in enumerate(lam):
if wl <= np.pi*a:
qext[idx_a, idx_wl] = A #regime A
elif np.pi*a < wl < 2*np.pi*a:
qext[idx_a, idx_wl] = q_c*(wl/lambda_c[idx_a])**(np.log10(q_c)/np.log10(2)) #regime B
elif wl >= 2*np.pi*a:
qext[idx_a, idx_wl] = q_c*(wl/lambda_c[idx_a])**(-2) #regime C
return qext
[docs]
def av_z(self, lam, nd, r, z):
"""Compute the vertical visual extinction profile.
Parameters
----------
lam : array_like
Wavelengths in microns.
nd : ndarray
Dust number density array of shape
``(n_sizes, len(r), len(z))`` in cm^-3.
r : array_like
Radial distances from the star in AU.
z : array_like
Normalized altitude coordinates (in units of scale height).
Returns
-------
avz : ndarray
Visual extinction (A_V) of shape ``(len(r), len(z))``,
in magnitudes.
"""
sizes = self.dust.sizes()
#extinction efficiency
lambda_c = 2*np.pi*sizes[-1]
qext = np.ones(( len(sizes[-1]), len(lam) ))
for idx_a, a in enumerate(sizes[-1]):
for idx_wl, wl in enumerate(lam):
if wl <= np.pi*a:
qext[idx_a, idx_wl] = 1 #regime A
elif np.pi*a < wl < 2*np.pi*a:
qext[idx_a, idx_wl] = self.params.q_c*(wl/lambda_c[idx_a])**(np.log10(self.params.q_c)/np.log10(2)) #regime B
elif wl >= 2*np.pi*a:
qext[idx_a, idx_wl] = self.params.q_c*(wl/lambda_c[idx_a])**(-2) #regime C
rr, zz = np.meshgrid(r, z, indexing='ij')
hg = self.scaleheight(rr)
zz = hg*zz
avz = np.zeros( (len(r), len(z)) )
for idx_r in range(0, len(r), 1):
for idx_z in range(1, len(z), 1):
for idx_a in range(0, len(sizes[-1]), 1):
avz[idx_r, idx_z] += 1.086*np.pi*nd[idx_a, idx_r, idx_z]*qext[idx_a, 12]*(zz[idx_r, idx_z-1] - zz[idx_r, idx_z])*(sizes[-1][idx_a]*1e-4)**2
avz[idx_r, :] = np.cumsum(avz[idx_r, :])
avz[:, 0] = avz[:, 1]
return avz
[docs]
def avz_ism(self, lam, A=2, q_c=4):
"""Compute visual extinction using the ISM H-column-density-to-Av conversion.
Parameters
----------
lam : array_like
Wavelengths in microns.
A : float, optional
Extinction efficiency in the geometric-optics regime (default 2).
q_c : float, optional
Extinction efficiency at the critical wavelength (default 4).
"""
# def av_z2(self ,mass, filelist, ):
# #--Av
# mass = d.grainmass() #g
# kext = np.zeros(len(mass))
# filelist = sorted(glob.glob('thermal/dustkap*'))
# for ai in range(0, len(mass)):
# kappa = pd.read_table(filelist[ai], sep="\s+", comment='#', header=None, skiprows=10)
# kext[ai] += (kappa[1].iloc[0] + kappa[2].iloc[0]) # k_extinction at wc ~ 550 nm.
# av_mid = np.ones(len(radii))
# for ri in range(len(radii)):
# z = m.grid.zchem*m.grid.hg_chem[0][ri]
# av_z = np.zeros(len(m.grid.zchem))
# av_z[0] = 0
# for zi in range(1, len(m.grid.zchem)):
# for ai in range(0, len(mass)):
# av_z[zi] += 1.086*m.grid.dustdensity_chem[0][ai, ri, zi]*mass[ai]*kext[ai]*(z[zi-1] - z[zi])
# av_z = np.cumsum(av_z)
# av_mid[ri] = av_z[-1]
# print(av_mid)
# #--
