Module diskchef.model.model

Expand source code
import pickle

from functools import cached_property
import pathlib
from typing import List, Type, Dict, Union
from dataclasses import dataclass, field
import warnings

import numpy as np
import astropy.wcs
import astropy.table
from astropy import units as u
import spectral_cube
from spectral_cube.utils import SpectralCubeWarning
from matplotlib import pyplot as plt

import diskchef.physics.base
from diskchef.chemistry import NonzeroChemistryWB2014
from diskchef.chemistry.scikit import SciKitChemistry
from diskchef.dust_opacity import dust_files
from diskchef import Line
from diskchef.maps import RadMCTherm, RadMCRTLines
from diskchef.physics import YorkeBodenheimer2008
from diskchef.physics.multidust import DustPopulation
from diskchef.physics.williams_best import WilliamsBest2014, WilliamsBest100au
from diskchef.engine.other import PathLike

from diskchef.uv import UVFits

warnings.filterwarnings(action='ignore', category=SpectralCubeWarning,
                        append=True)


@dataclass
class ModelFit:
    disk: str
    directory: Union[pathlib.Path, str] = None
    line_list: List[Line] = None
    physics_class: Type[diskchef.physics.base.PhysicsModel] = WilliamsBest100au
    physics_params: dict = field(default_factory=dict)
    chemistry_class: Type[diskchef.chemistry.base.ChemistryBase] = NonzeroChemistryWB2014
    chemical_params: dict = field(default_factory=dict)
    mstar: u.Msun = 1 * u.Msun
    rstar: u.au = None
    tstar: u.K = None
    inclination: u.deg = 0 * u.deg
    position_angle: u.deg = 0 * u.deg
    distance: u.pc = 100 * u.pc
    nphot_therm: int = 1e7
    npix: int = 100
    channels: int = 31
    dust_opacity_file: Union[pathlib.Path, str] = dust_files("diana")[0]
    radial_bins_rt: int = None
    vertical_bins_rt: int = None
    line_window_width: u.km / u.s = 15 * u.km / u.s
    radmc_lines_run_kwargs: dict = field(default_factory=dict)
    mctherm_threads = 4

    def __post_init__(self):
        self.dust = None
        self.disk_physical_model = None
        self.disk_chemical_model = None
        self.chi2_dict = None
        self._check_radius_temperature()
        self._update_defaults()
        self._initialize_working_dir()
        self.initialize_physics()
        self.initialize_dust()
        self.initialize_chemistry()

    def run_chemistry(self):
        self.disk_chemical_model.run_chemistry()
        self.add_co_isotopologs()

    def run_line_radiative_transfer(
            self,
            run_kwargs: dict = None,
            **kwargs
    ):
        """Run line radiative transfer"""
        folder_rt_gas = self.gas_directory
        folder_rt_gas.mkdir(parents=True, exist_ok=True)

        disk_map = RadMCRTLines(
            chemistry=self.disk_chemical_model, line_list=self.line_list,
            radii_bins=self.radial_bins_rt, theta_bins=self.vertical_bins_rt,
            folder=folder_rt_gas, **kwargs
        )

        disk_map.create_files(channels_per_line=self.channels, window_width=self.line_window_width)

        disk_map.run(
            inclination=self.inclination,
            position_angle=self.position_angle,
            distance=self.distance,
            npix=self.npix,
            **self.radmc_lines_run_kwargs
        )

    def add_co_isotopologs(self, _13c: float = 77, _18o: float = 560):
        self.disk_chemical_model.table['13CO'] = self.disk_chemical_model.table['CO'] / _13c
        self.disk_chemical_model.table['C18O'] = self.disk_chemical_model.table['CO'] / _18o
        self.disk_chemical_model.table['13C18O'] = self.disk_chemical_model.table['CO'] / (_13c * _18o)
        # remember to fix the next abundance
        # self.disk_chemical_model.table['C17O'] = self.disk_chemical_model.table['CO'] / (
        #         560 * 5)  # http://articles.adsabs.harvard.edu/pdf/1994ARA%26A..32..191W

    def initialize_chemistry(self):
        self.disk_chemical_model = self.chemistry_class(physics=self.disk_physical_model, **self.chemical_params)

    def initialize_dust(self):
        self.dust = DustPopulation(self.dust_opacity_file,
                                   table=self.disk_physical_model.table,
                                   name="Dust")
        self.dust.write_to_table()

    def initialize_physics(self):
        self.disk_physical_model = self.physics_class(**self.physics_params)

    def _update_defaults(self):
        if self.radial_bins_rt is None:
            self.radial_bins_rt = self.physics_params.get("radial_bins", 100)
        if self.vertical_bins_rt is None:
            self.radial_bins_rt = self.physics_params.get("vertical_bins", 100)

    def _initialize_working_dir(self):
        if self.directory is None:
            self.directory = Path(self.disk)
        else:
            self.directory = pathlib.Path(self.directory)
        self.directory.mkdir(exist_ok=True, parents=True)
        with open(self.directory / "model_description.txt", "w") as fff:
            fff.write(repr(self))
            fff.write("\n")
            fff.write(str(self.__dict__))

    def _check_radius_temperature(self):
        if self.rstar is None and self.tstar is None:
            yb = YorkeBodenheimer2008()
            self.rstar = yb.radius(self.mstar)
            self.tstar = yb.effective_temperature(self.mstar)

    def mctherm(self, threads=None):
        """Run thermal radiative transfer for dust temperature calculation"""
        if threads is None:
            threads = self.mctherm_threads
        folder_rt_dust = self.dust_directory
        folder_rt_dust.mkdir(parents=True, exist_ok=True)

        self.mctherm_model = RadMCTherm(
            chemistry=self.disk_chemical_model,
            folder=folder_rt_dust,
            nphot_therm=self.nphot_therm,
            star_radius=self.rstar,
            star_effective_temperature=self.tstar
        )

        self.mctherm_model.create_files()
        self.mctherm_model.run(threads=threads)
        self.mctherm_model.read_dust_temperature()

    @property
    def gas_directory(self):
        return self.directory / "radmc_gas"

    @property
    def dust_directory(self):
        return self.directory / "radmc_dust"

    def chi2(self, uvs: Dict[str, UVFits]):
        """

        Args:
            uvs: dictionary in a form of {line.name: uvt for line is self.line_list}

        Returns:
            sum of chi2 between uvs and lines
        """

        self.chi2_dict = {
            name: uv.chi2_with(self.gas_directory / f"{name}_image.fits")
            for name, uv in uvs.items()
            if name in [line.name for line in self.line_list]
        }
        return sum(self.chi2_dict.values())


@dataclass
class Model:
    """Class to run a simulation"""
    disk: str
    rstar: u.au
    tstar: u.K
    params: dict
    inc: u.deg
    PA: u.deg
    distance: u.pc
    nphot_therm: int = 1e7
    run: bool = True
    line_list: List[Line] = None
    npix: int = 100
    channels: int = 31
    run_mctherm: bool = True
    radial_bins_rt: int = None
    vertical_bins_rt: int = None
    folder: pathlib.Path = None
    physics_class: Type[diskchef.physics.base.PhysicsModel] = WilliamsBest2014
    scikit_model: PathLike = None

    def __post_init__(self):
        warnings.warn("Model class is deprecated in favor of ModelFit, which is being developed!", DeprecationWarning)
        if self.folder is None:
            self.folder = pathlib.Path(self.disk)
        else:
            self.folder = pathlib.Path(self.folder)
        self.folder.mkdir(exist_ok=True, parents=True)
        with open(self.folder / "model_description.txt", "w") as fff:
            fff.write(repr(self))
            fff.write("\n")
            fff.write(str(self.__dict__))

        if self.radial_bins_rt is None:
            self.radial_bins_rt = self.params.get("radial_bins", 100)
        if self.vertical_bins_rt is None:
            self.radial_bins_rt = self.params.get("vertical_bins", 100)

        self.disk_physical_model = self.physics_class(**self.params)
        dust = DustPopulation(dust_files("diana")[0],
                              table=self.disk_physical_model.table,
                              name="DIANA dust")
        dust.write_to_table()
        self.disk_chemical_model = SciKitChemistry(self.disk_physical_model, model=self.scikit_model)
        if self.run:
            self.run_simulation()

    def run_simulation(self):
        """Run all the steps of the simulation"""
        if self.run_mctherm: self.mctherm()
        self.chemistry()
        self.plot()
        self.image()
        self.photometry()

    def mctherm(self):
        """Run thermal radiative transfer for dust temperature calculation"""
        folder_rt_dust = self.folder / "radmc_dust"
        # folder_rt_dust.mkdir(parents=True, exist_ok=True)

        self.mctherm_model = RadMCTherm(
            chemistry=self.disk_chemical_model,
            folder=folder_rt_dust,
            nphot_therm=self.nphot_therm,
            star_radius=self.rstar,
            star_effective_temperature=self.tstar
        )

        self.mctherm_model.create_files()
        self.mctherm_model.run(threads=4)
        self.mctherm_model.read_dust_temperature()

    def image(
            self,
            radii_bins: int = 100, theta_bins: int = 100,
            run_kwargs: dict = None, window_width: u.km / u.s = 15 * u.km / u.s,
            **kwargs
    ):
        """Run line radiative transfer"""
        folder_rt_gas = self.folder / "radmc_gas"
        # folder_rt_gas.mkdir(parents=True, exist_ok=True)

        disk_map = RadMCRTLines(
            chemistry=self.disk_chemical_model, line_list=self.line_list,
            radii_bins=radii_bins, theta_bins=theta_bins,
            folder=folder_rt_gas, **kwargs
        )

        disk_map.create_files(channels_per_line=self.channels, window_width=window_width)

        if run_kwargs is None: run_kwargs = {}
        disk_map.run(inclination=self.inc,
                     position_angle=self.PA,
                     distance=self.distance,
                     npix=self.npix,
                     **run_kwargs
                     )

    def chemistry(self):
        """Run chemistry"""
        self.disk_chemical_model.run_chemistry()

        self.disk_chemical_model.table['13CO'] = self.disk_chemical_model.table['CO'] / 77
        self.disk_chemical_model.table['C18O'] = self.disk_chemical_model.table['CO'] / 560
        # remember to fix the next abundance
        self.disk_chemical_model.table['C17O'] = self.disk_chemical_model.table['CO'] / (
                560 * 5)  # http://articles.adsabs.harvard.edu/pdf/1994ARA%26A..32..191W
        self.disk_chemical_model.table['13C18O'] = self.disk_chemical_model.table['CO'] / (77 * 560)

    def plot(self, save=None, **kwargs):
        """Plot physical and chemical structure"""
        fig, ax = plt.subplots(2, 2, sharex='all', sharey='all', figsize=(11, 10))

        self.disk_physical_model.plot_density(axes=ax[0, 0])
        tempplot = self.disk_physical_model.plot_temperatures(axes=ax[1, 0])
        tempplot.contours("Gas temperature", [20, 40] * u.K)

        self.disk_chemical_model.plot_chemistry("CO", "CO", axes=ax[0, 1])
        coplot = self.disk_chemical_model.plot_absolute_chemistry("CO", "CO", axes=ax[1, 1], maxdepth=1e9)
        coplot.contours("Gas temperature", [20, 40] * u.K, clabel_kwargs={"fmt": "T=%d K"})
        if save:
            fig.savefig(self.folder / "structure.png")
            fig.savefig(self.folder / "structure.pdf")
        return fig

    def pickle(self):
        """Pickle the model class in a file -- not working"""
        with open(self.folder / "model.pkl", "wb") as pkl:
            pickle.dump(self, pkl)

    @classmethod
    def from_picke(cls, path: pathlib.Path):
        """Return the model class from the pickle file  -- experimental"""
        with open(path, "rb") as pkl:
            return pickle.load(pkl)

    @cached_property
    def output_fluxes(self) -> pathlib.Path:
        """Path to output `fluxes.ecsv` table"""
        return self.folder / "fluxes.ecsv"

    def photometry(self):
        fluxes = []
        transitions = [line.name for line in self.line_list]
        for transition in transitions:
            file = self.folder / "radmc_gas" / f"{transition}_image.fits"
            scube = spectral_cube.SpectralCube.read(file)
            scube = scube.with_spectral_unit(u.km / u.s, velocity_convention='radio')
            pixel_area_units = u.Unit(scube.wcs.celestial.world_axis_units[0]) * u.Unit(
                scube.wcs.celestial.world_axis_units[1])
            pixel_area = astropy.wcs.utils.proj_plane_pixel_area(scube.wcs.celestial) * pixel_area_units

            spectrum = (scube * pixel_area).to(u.mJy).sum(axis=(1, 2))  # 1d spectrum in Jy
            flux = np.abs(np.trapz(spectrum, scube.spectral_axis))
            tbl = astropy.table.QTable(
                [scube.spectral_axis, u.Quantity(spectrum)],
                meta={"flux": flux},
                names=["Velocity", "Flux density"],
            )
            tbl.write(self.folder / f"{transition}_spectrum.ecsv", overwrite=True)
            fluxes.append(flux)
        fluxes_tbl = astropy.table.QTable(
            [transitions, fluxes], names=["Transition", "Flux"], meta={"Source": "Default"}
        )
        fluxes_tbl.write(self.output_fluxes, overwrite=True)
        self.fluxes_tbl = fluxes_tbl

Classes

class Model (disk: str, rstar: Unit("AU"), tstar: Unit("K"), params: dict, inc: Unit("deg"), PA: Unit("deg"), distance: Unit("pc"), nphot_therm: int = 10000000.0, run: bool = True, line_list: List[Line] = None, npix: int = 100, channels: int = 31, run_mctherm: bool = True, radial_bins_rt: int = None, vertical_bins_rt: int = None, folder: pathlib.Path = None, physics_class: Type[PhysicsModel] = diskchef.physics.williams_best.WilliamsBest2014, scikit_model: Union[str, os.PathLike] = None)

Class to run a simulation

Expand source code
@dataclass
class Model:
    """Class to run a simulation"""
    disk: str
    rstar: u.au
    tstar: u.K
    params: dict
    inc: u.deg
    PA: u.deg
    distance: u.pc
    nphot_therm: int = 1e7
    run: bool = True
    line_list: List[Line] = None
    npix: int = 100
    channels: int = 31
    run_mctherm: bool = True
    radial_bins_rt: int = None
    vertical_bins_rt: int = None
    folder: pathlib.Path = None
    physics_class: Type[diskchef.physics.base.PhysicsModel] = WilliamsBest2014
    scikit_model: PathLike = None

    def __post_init__(self):
        warnings.warn("Model class is deprecated in favor of ModelFit, which is being developed!", DeprecationWarning)
        if self.folder is None:
            self.folder = pathlib.Path(self.disk)
        else:
            self.folder = pathlib.Path(self.folder)
        self.folder.mkdir(exist_ok=True, parents=True)
        with open(self.folder / "model_description.txt", "w") as fff:
            fff.write(repr(self))
            fff.write("\n")
            fff.write(str(self.__dict__))

        if self.radial_bins_rt is None:
            self.radial_bins_rt = self.params.get("radial_bins", 100)
        if self.vertical_bins_rt is None:
            self.radial_bins_rt = self.params.get("vertical_bins", 100)

        self.disk_physical_model = self.physics_class(**self.params)
        dust = DustPopulation(dust_files("diana")[0],
                              table=self.disk_physical_model.table,
                              name="DIANA dust")
        dust.write_to_table()
        self.disk_chemical_model = SciKitChemistry(self.disk_physical_model, model=self.scikit_model)
        if self.run:
            self.run_simulation()

    def run_simulation(self):
        """Run all the steps of the simulation"""
        if self.run_mctherm: self.mctherm()
        self.chemistry()
        self.plot()
        self.image()
        self.photometry()

    def mctherm(self):
        """Run thermal radiative transfer for dust temperature calculation"""
        folder_rt_dust = self.folder / "radmc_dust"
        # folder_rt_dust.mkdir(parents=True, exist_ok=True)

        self.mctherm_model = RadMCTherm(
            chemistry=self.disk_chemical_model,
            folder=folder_rt_dust,
            nphot_therm=self.nphot_therm,
            star_radius=self.rstar,
            star_effective_temperature=self.tstar
        )

        self.mctherm_model.create_files()
        self.mctherm_model.run(threads=4)
        self.mctherm_model.read_dust_temperature()

    def image(
            self,
            radii_bins: int = 100, theta_bins: int = 100,
            run_kwargs: dict = None, window_width: u.km / u.s = 15 * u.km / u.s,
            **kwargs
    ):
        """Run line radiative transfer"""
        folder_rt_gas = self.folder / "radmc_gas"
        # folder_rt_gas.mkdir(parents=True, exist_ok=True)

        disk_map = RadMCRTLines(
            chemistry=self.disk_chemical_model, line_list=self.line_list,
            radii_bins=radii_bins, theta_bins=theta_bins,
            folder=folder_rt_gas, **kwargs
        )

        disk_map.create_files(channels_per_line=self.channels, window_width=window_width)

        if run_kwargs is None: run_kwargs = {}
        disk_map.run(inclination=self.inc,
                     position_angle=self.PA,
                     distance=self.distance,
                     npix=self.npix,
                     **run_kwargs
                     )

    def chemistry(self):
        """Run chemistry"""
        self.disk_chemical_model.run_chemistry()

        self.disk_chemical_model.table['13CO'] = self.disk_chemical_model.table['CO'] / 77
        self.disk_chemical_model.table['C18O'] = self.disk_chemical_model.table['CO'] / 560
        # remember to fix the next abundance
        self.disk_chemical_model.table['C17O'] = self.disk_chemical_model.table['CO'] / (
                560 * 5)  # http://articles.adsabs.harvard.edu/pdf/1994ARA%26A..32..191W
        self.disk_chemical_model.table['13C18O'] = self.disk_chemical_model.table['CO'] / (77 * 560)

    def plot(self, save=None, **kwargs):
        """Plot physical and chemical structure"""
        fig, ax = plt.subplots(2, 2, sharex='all', sharey='all', figsize=(11, 10))

        self.disk_physical_model.plot_density(axes=ax[0, 0])
        tempplot = self.disk_physical_model.plot_temperatures(axes=ax[1, 0])
        tempplot.contours("Gas temperature", [20, 40] * u.K)

        self.disk_chemical_model.plot_chemistry("CO", "CO", axes=ax[0, 1])
        coplot = self.disk_chemical_model.plot_absolute_chemistry("CO", "CO", axes=ax[1, 1], maxdepth=1e9)
        coplot.contours("Gas temperature", [20, 40] * u.K, clabel_kwargs={"fmt": "T=%d K"})
        if save:
            fig.savefig(self.folder / "structure.png")
            fig.savefig(self.folder / "structure.pdf")
        return fig

    def pickle(self):
        """Pickle the model class in a file -- not working"""
        with open(self.folder / "model.pkl", "wb") as pkl:
            pickle.dump(self, pkl)

    @classmethod
    def from_picke(cls, path: pathlib.Path):
        """Return the model class from the pickle file  -- experimental"""
        with open(path, "rb") as pkl:
            return pickle.load(pkl)

    @cached_property
    def output_fluxes(self) -> pathlib.Path:
        """Path to output `fluxes.ecsv` table"""
        return self.folder / "fluxes.ecsv"

    def photometry(self):
        fluxes = []
        transitions = [line.name for line in self.line_list]
        for transition in transitions:
            file = self.folder / "radmc_gas" / f"{transition}_image.fits"
            scube = spectral_cube.SpectralCube.read(file)
            scube = scube.with_spectral_unit(u.km / u.s, velocity_convention='radio')
            pixel_area_units = u.Unit(scube.wcs.celestial.world_axis_units[0]) * u.Unit(
                scube.wcs.celestial.world_axis_units[1])
            pixel_area = astropy.wcs.utils.proj_plane_pixel_area(scube.wcs.celestial) * pixel_area_units

            spectrum = (scube * pixel_area).to(u.mJy).sum(axis=(1, 2))  # 1d spectrum in Jy
            flux = np.abs(np.trapz(spectrum, scube.spectral_axis))
            tbl = astropy.table.QTable(
                [scube.spectral_axis, u.Quantity(spectrum)],
                meta={"flux": flux},
                names=["Velocity", "Flux density"],
            )
            tbl.write(self.folder / f"{transition}_spectrum.ecsv", overwrite=True)
            fluxes.append(flux)
        fluxes_tbl = astropy.table.QTable(
            [transitions, fluxes], names=["Transition", "Flux"], meta={"Source": "Default"}
        )
        fluxes_tbl.write(self.output_fluxes, overwrite=True)
        self.fluxes_tbl = fluxes_tbl

Class variables

var PA : Unit("deg")
var channels : int
var disk : str
var distance : Unit("pc")
var folder : pathlib.Path
var inc : Unit("deg")
var line_list : List[Line]
var nphot_therm : int
var npix : int
var params : dict
var physics_class : Type[PhysicsModel]

Class that defines physical model as in Williams & Best 2014 https://iopscience.iop.org/article/10.1088/0004-637X/788/1/59/pdf

Usage:

>>> # Initializing physics class with default parameters
>>> physics = WilliamsBest2014()
>>> physics  # doctest: +NORMALIZE_WHITESPACE
WilliamsBest2014(star_mass=<Quantity 1. solMass>, xray_plasma_temperature=<Quantity 10000000. K>,
                 xray_luminosity=<Quantity 1.e+31 erg / s>, cr_padovani_use_l=False, r_min=<Quantity 0.1 AU>,
                 r_max=<Quantity 500. AU>, zr_max=0.7, radial_bins=100, vertical_bins=100, dust_to_gas=0.01,
                 gas_mass=<Quantity 0.001 solMass>, tapering_radius=<Quantity 100. AU>,
                 tapering_gamma=0.75, midplane_temperature_1au=<Quantity 200. K>,
                  atmosphere_temperature_1au=<Quantity 1000. K>, temperature_slope=0.55,
                  molar_mass=<Quantity 2.33 g / mol>, inner_radius=<Quantity 1. AU>, inner_depletion=1e-06)
>>> # Defaults can be overridden
>>> WilliamsBest2014(star_mass=2 * u.solMass, r_max=200 * u.au)  # doctest: +NORMALIZE_WHITESPACE
WilliamsBest2014(star_mass=<Quantity 2. solMass>, xray_plasma_temperature=<Quantity 10000000. K>,
                 xray_luminosity=<Quantity 1.e+31 erg / s>, cr_padovani_use_l=False, r_min=<Quantity 0.1 AU>,
                 r_max=<Quantity 200. AU>, zr_max=0.7, radial_bins=100, vertical_bins=100, dust_to_gas=0.01,
                 gas_mass=<Quantity 0.001 solMass>, tapering_radius=<Quantity 100. AU>, tapering_gamma=0.75,
                 midplane_temperature_1au=<Quantity 200. K>, atmosphere_temperature_1au=<Quantity 1000. K>,
                 temperature_slope=0.55, molar_mass=<Quantity 2.33 g / mol>, inner_radius=<Quantity 1. AU>,
                 inner_depletion=1e-06)
>>> # Generate physics on 3x3 grid
>>> physics_small_grid = WilliamsBest2014(vertical_bins=3, radial_bins=3)
>>> # table attribute stores the table with the model. Called first time, calculates the table
>>> table = physics_small_grid.table
>>> # print table with floats in exponential format
>>> table  # doctest: +NORMALIZE_WHITESPACE
<CTable length=9>
   Radius       Height    Height to radius Gas density  Dust density Gas temperature Dust temperature
     AU           AU                         g / cm3      g / cm3           K               K
  float64      float64        float64        float64      float64        float64         float64
------------ ------------ ---------------- ------------ ------------ --------------- ----------------
1.000000e-01 0.000000e+00     0.000000e+00 1.153290e-15 1.153290e-17    7.096268e+02     7.096268e+02
1.000000e-01 3.500000e-02     3.500000e-01 5.024587e-34 5.024587e-36    3.548134e+03     3.548134e+03
1.000000e-01 7.000000e-02     7.000000e-01 5.921268e-72 5.921268e-74    3.548134e+03     3.548134e+03
7.071068e+00 0.000000e+00     0.000000e+00 2.285168e-13 2.285168e-15    6.820453e+01     6.820453e+01
7.071068e+00 2.474874e+00     3.500000e-01 2.716386e-17 2.716386e-19    3.410227e+02     3.410227e+02
7.071068e+00 4.949747e+00     7.000000e-01 7.189026e-23 7.189026e-25    3.410227e+02     3.410227e+02
5.000000e+02 0.000000e+00     0.000000e+00 2.871710e-20 2.871710e-22    6.555359e+00     6.555359e+00
5.000000e+02 1.750000e+02     3.500000e-01 6.929036e-22 6.929036e-24    2.625083e+01     2.625083e+01
5.000000e+02 3.500000e+02     7.000000e-01 8.170464e-23 8.170464e-25    3.277680e+01     3.277680e+01
>>> physics_wrong_unit = WilliamsBest2014(star_mass=2)  # Does NOT raise an exception
>>> physics_wrong_unit.table  # but will cause unit conversion errors later
Traceback (most recent call last):
   ...
astropy.units.core.UnitConversionError: 'AU(3/2) J(1/2) kg(1/2) s / (g(1/2) m(3/2))' and 'AU' (length) are not convertible
var radial_bins_rt : int
var rstar : Unit("AU")
var run : bool
var run_mctherm : bool
var scikit_model : Union[str, os.PathLike]
var tstar : Unit("K")
var vertical_bins_rt : int

Static methods

def from_picke(path: pathlib.Path)

Return the model class from the pickle file – experimental

Expand source code
@classmethod
def from_picke(cls, path: pathlib.Path):
    """Return the model class from the pickle file  -- experimental"""
    with open(path, "rb") as pkl:
        return pickle.load(pkl)

Instance variables

var output_fluxes

Path to output fluxes.ecsv table

Expand source code
def __get__(self, instance, owner=None):
    if instance is None:
        return self
    if self.attrname is None:
        raise TypeError(
            "Cannot use cached_property instance without calling __set_name__ on it.")
    try:
        cache = instance.__dict__
    except AttributeError:  # not all objects have __dict__ (e.g. class defines slots)
        msg = (
            f"No '__dict__' attribute on {type(instance).__name__!r} "
            f"instance to cache {self.attrname!r} property."
        )
        raise TypeError(msg) from None
    val = cache.get(self.attrname, _NOT_FOUND)
    if val is _NOT_FOUND:
        with self.lock:
            # check if another thread filled cache while we awaited lock
            val = cache.get(self.attrname, _NOT_FOUND)
            if val is _NOT_FOUND:
                val = self.func(instance)
                try:
                    cache[self.attrname] = val
                except TypeError:
                    msg = (
                        f"The '__dict__' attribute on {type(instance).__name__!r} instance "
                        f"does not support item assignment for caching {self.attrname!r} property."
                    )
                    raise TypeError(msg) from None
    return val

Methods

def chemistry(self)

Run chemistry

Expand source code
def chemistry(self):
    """Run chemistry"""
    self.disk_chemical_model.run_chemistry()

    self.disk_chemical_model.table['13CO'] = self.disk_chemical_model.table['CO'] / 77
    self.disk_chemical_model.table['C18O'] = self.disk_chemical_model.table['CO'] / 560
    # remember to fix the next abundance
    self.disk_chemical_model.table['C17O'] = self.disk_chemical_model.table['CO'] / (
            560 * 5)  # http://articles.adsabs.harvard.edu/pdf/1994ARA%26A..32..191W
    self.disk_chemical_model.table['13C18O'] = self.disk_chemical_model.table['CO'] / (77 * 560)
def image(self, radii_bins: int = 100, theta_bins: int = 100, run_kwargs: dict = None, window_width: Unit("km / s") = <Quantity 15. km / s>, **kwargs)

Run line radiative transfer

Expand source code
def image(
        self,
        radii_bins: int = 100, theta_bins: int = 100,
        run_kwargs: dict = None, window_width: u.km / u.s = 15 * u.km / u.s,
        **kwargs
):
    """Run line radiative transfer"""
    folder_rt_gas = self.folder / "radmc_gas"
    # folder_rt_gas.mkdir(parents=True, exist_ok=True)

    disk_map = RadMCRTLines(
        chemistry=self.disk_chemical_model, line_list=self.line_list,
        radii_bins=radii_bins, theta_bins=theta_bins,
        folder=folder_rt_gas, **kwargs
    )

    disk_map.create_files(channels_per_line=self.channels, window_width=window_width)

    if run_kwargs is None: run_kwargs = {}
    disk_map.run(inclination=self.inc,
                 position_angle=self.PA,
                 distance=self.distance,
                 npix=self.npix,
                 **run_kwargs
                 )
def mctherm(self)

Run thermal radiative transfer for dust temperature calculation

Expand source code
def mctherm(self):
    """Run thermal radiative transfer for dust temperature calculation"""
    folder_rt_dust = self.folder / "radmc_dust"
    # folder_rt_dust.mkdir(parents=True, exist_ok=True)

    self.mctherm_model = RadMCTherm(
        chemistry=self.disk_chemical_model,
        folder=folder_rt_dust,
        nphot_therm=self.nphot_therm,
        star_radius=self.rstar,
        star_effective_temperature=self.tstar
    )

    self.mctherm_model.create_files()
    self.mctherm_model.run(threads=4)
    self.mctherm_model.read_dust_temperature()
def photometry(self)
Expand source code
def photometry(self):
    fluxes = []
    transitions = [line.name for line in self.line_list]
    for transition in transitions:
        file = self.folder / "radmc_gas" / f"{transition}_image.fits"
        scube = spectral_cube.SpectralCube.read(file)
        scube = scube.with_spectral_unit(u.km / u.s, velocity_convention='radio')
        pixel_area_units = u.Unit(scube.wcs.celestial.world_axis_units[0]) * u.Unit(
            scube.wcs.celestial.world_axis_units[1])
        pixel_area = astropy.wcs.utils.proj_plane_pixel_area(scube.wcs.celestial) * pixel_area_units

        spectrum = (scube * pixel_area).to(u.mJy).sum(axis=(1, 2))  # 1d spectrum in Jy
        flux = np.abs(np.trapz(spectrum, scube.spectral_axis))
        tbl = astropy.table.QTable(
            [scube.spectral_axis, u.Quantity(spectrum)],
            meta={"flux": flux},
            names=["Velocity", "Flux density"],
        )
        tbl.write(self.folder / f"{transition}_spectrum.ecsv", overwrite=True)
        fluxes.append(flux)
    fluxes_tbl = astropy.table.QTable(
        [transitions, fluxes], names=["Transition", "Flux"], meta={"Source": "Default"}
    )
    fluxes_tbl.write(self.output_fluxes, overwrite=True)
    self.fluxes_tbl = fluxes_tbl
def pickle(self)

Pickle the model class in a file – not working

Expand source code
def pickle(self):
    """Pickle the model class in a file -- not working"""
    with open(self.folder / "model.pkl", "wb") as pkl:
        pickle.dump(self, pkl)
def plot(self, save=None, **kwargs)

Plot physical and chemical structure

Expand source code
def plot(self, save=None, **kwargs):
    """Plot physical and chemical structure"""
    fig, ax = plt.subplots(2, 2, sharex='all', sharey='all', figsize=(11, 10))

    self.disk_physical_model.plot_density(axes=ax[0, 0])
    tempplot = self.disk_physical_model.plot_temperatures(axes=ax[1, 0])
    tempplot.contours("Gas temperature", [20, 40] * u.K)

    self.disk_chemical_model.plot_chemistry("CO", "CO", axes=ax[0, 1])
    coplot = self.disk_chemical_model.plot_absolute_chemistry("CO", "CO", axes=ax[1, 1], maxdepth=1e9)
    coplot.contours("Gas temperature", [20, 40] * u.K, clabel_kwargs={"fmt": "T=%d K"})
    if save:
        fig.savefig(self.folder / "structure.png")
        fig.savefig(self.folder / "structure.pdf")
    return fig
def run_simulation(self)

Run all the steps of the simulation

Expand source code
def run_simulation(self):
    """Run all the steps of the simulation"""
    if self.run_mctherm: self.mctherm()
    self.chemistry()
    self.plot()
    self.image()
    self.photometry()
class ModelFit (disk: str, directory: Union[str, pathlib.Path] = None, line_list: List[Line] = None, physics_class: Type[PhysicsModel] = diskchef.physics.williams_best.WilliamsBest100au, physics_params: dict = <factory>, chemistry_class: Type[ChemistryBase] = diskchef.chemistry.williams_best_co.NonzeroChemistryWB2014, chemical_params: dict = <factory>, mstar: Unit("solMass") = <Quantity 1. solMass>, rstar: Unit("AU") = None, tstar: Unit("K") = None, inclination: Unit("deg") = <Quantity 0. deg>, position_angle: Unit("deg") = <Quantity 0. deg>, distance: Unit("pc") = <Quantity 100. pc>, nphot_therm: int = 10000000.0, npix: int = 100, channels: int = 31, dust_opacity_file: Union[str, pathlib.Path] = '/builds/SmirnGreg/diskchef/diskchef/dust_opacity/files/dustkapscatmat_diana.inp', radial_bins_rt: int = None, vertical_bins_rt: int = None, line_window_width: Unit("km / s") = <Quantity 15. km / s>, radmc_lines_run_kwargs: dict = <factory>)

ModelFit(disk: str, directory: Union[pathlib.Path, str] = None, line_list: List[diskchef.lamda.line.Line] = None, physics_class: Type[diskchef.physics.base.PhysicsModel] = , physics_params: dict = , chemistry_class: Type[diskchef.chemistry.base.ChemistryBase] = , chemical_params: dict = , mstar: Unit("solMass") = , rstar: Unit("AU") = None, tstar: Unit("K") = None, inclination: Unit("deg") = , position_angle: Unit("deg") = , distance: Unit("pc") = , nphot_therm: int = 10000000.0, npix: int = 100, channels: int = 31, dust_opacity_file: Union[pathlib.Path, str] = '/builds/SmirnGreg/diskchef/diskchef/dust_opacity/files/dustkapscatmat_diana.inp', radial_bins_rt: int = None, vertical_bins_rt: int = None, line_window_width: Unit("km / s") = , radmc_lines_run_kwargs: dict = )

Expand source code
@dataclass
class ModelFit:
    disk: str
    directory: Union[pathlib.Path, str] = None
    line_list: List[Line] = None
    physics_class: Type[diskchef.physics.base.PhysicsModel] = WilliamsBest100au
    physics_params: dict = field(default_factory=dict)
    chemistry_class: Type[diskchef.chemistry.base.ChemistryBase] = NonzeroChemistryWB2014
    chemical_params: dict = field(default_factory=dict)
    mstar: u.Msun = 1 * u.Msun
    rstar: u.au = None
    tstar: u.K = None
    inclination: u.deg = 0 * u.deg
    position_angle: u.deg = 0 * u.deg
    distance: u.pc = 100 * u.pc
    nphot_therm: int = 1e7
    npix: int = 100
    channels: int = 31
    dust_opacity_file: Union[pathlib.Path, str] = dust_files("diana")[0]
    radial_bins_rt: int = None
    vertical_bins_rt: int = None
    line_window_width: u.km / u.s = 15 * u.km / u.s
    radmc_lines_run_kwargs: dict = field(default_factory=dict)
    mctherm_threads = 4

    def __post_init__(self):
        self.dust = None
        self.disk_physical_model = None
        self.disk_chemical_model = None
        self.chi2_dict = None
        self._check_radius_temperature()
        self._update_defaults()
        self._initialize_working_dir()
        self.initialize_physics()
        self.initialize_dust()
        self.initialize_chemistry()

    def run_chemistry(self):
        self.disk_chemical_model.run_chemistry()
        self.add_co_isotopologs()

    def run_line_radiative_transfer(
            self,
            run_kwargs: dict = None,
            **kwargs
    ):
        """Run line radiative transfer"""
        folder_rt_gas = self.gas_directory
        folder_rt_gas.mkdir(parents=True, exist_ok=True)

        disk_map = RadMCRTLines(
            chemistry=self.disk_chemical_model, line_list=self.line_list,
            radii_bins=self.radial_bins_rt, theta_bins=self.vertical_bins_rt,
            folder=folder_rt_gas, **kwargs
        )

        disk_map.create_files(channels_per_line=self.channels, window_width=self.line_window_width)

        disk_map.run(
            inclination=self.inclination,
            position_angle=self.position_angle,
            distance=self.distance,
            npix=self.npix,
            **self.radmc_lines_run_kwargs
        )

    def add_co_isotopologs(self, _13c: float = 77, _18o: float = 560):
        self.disk_chemical_model.table['13CO'] = self.disk_chemical_model.table['CO'] / _13c
        self.disk_chemical_model.table['C18O'] = self.disk_chemical_model.table['CO'] / _18o
        self.disk_chemical_model.table['13C18O'] = self.disk_chemical_model.table['CO'] / (_13c * _18o)
        # remember to fix the next abundance
        # self.disk_chemical_model.table['C17O'] = self.disk_chemical_model.table['CO'] / (
        #         560 * 5)  # http://articles.adsabs.harvard.edu/pdf/1994ARA%26A..32..191W

    def initialize_chemistry(self):
        self.disk_chemical_model = self.chemistry_class(physics=self.disk_physical_model, **self.chemical_params)

    def initialize_dust(self):
        self.dust = DustPopulation(self.dust_opacity_file,
                                   table=self.disk_physical_model.table,
                                   name="Dust")
        self.dust.write_to_table()

    def initialize_physics(self):
        self.disk_physical_model = self.physics_class(**self.physics_params)

    def _update_defaults(self):
        if self.radial_bins_rt is None:
            self.radial_bins_rt = self.physics_params.get("radial_bins", 100)
        if self.vertical_bins_rt is None:
            self.radial_bins_rt = self.physics_params.get("vertical_bins", 100)

    def _initialize_working_dir(self):
        if self.directory is None:
            self.directory = Path(self.disk)
        else:
            self.directory = pathlib.Path(self.directory)
        self.directory.mkdir(exist_ok=True, parents=True)
        with open(self.directory / "model_description.txt", "w") as fff:
            fff.write(repr(self))
            fff.write("\n")
            fff.write(str(self.__dict__))

    def _check_radius_temperature(self):
        if self.rstar is None and self.tstar is None:
            yb = YorkeBodenheimer2008()
            self.rstar = yb.radius(self.mstar)
            self.tstar = yb.effective_temperature(self.mstar)

    def mctherm(self, threads=None):
        """Run thermal radiative transfer for dust temperature calculation"""
        if threads is None:
            threads = self.mctherm_threads
        folder_rt_dust = self.dust_directory
        folder_rt_dust.mkdir(parents=True, exist_ok=True)

        self.mctherm_model = RadMCTherm(
            chemistry=self.disk_chemical_model,
            folder=folder_rt_dust,
            nphot_therm=self.nphot_therm,
            star_radius=self.rstar,
            star_effective_temperature=self.tstar
        )

        self.mctherm_model.create_files()
        self.mctherm_model.run(threads=threads)
        self.mctherm_model.read_dust_temperature()

    @property
    def gas_directory(self):
        return self.directory / "radmc_gas"

    @property
    def dust_directory(self):
        return self.directory / "radmc_dust"

    def chi2(self, uvs: Dict[str, UVFits]):
        """

        Args:
            uvs: dictionary in a form of {line.name: uvt for line is self.line_list}

        Returns:
            sum of chi2 between uvs and lines
        """

        self.chi2_dict = {
            name: uv.chi2_with(self.gas_directory / f"{name}_image.fits")
            for name, uv in uvs.items()
            if name in [line.name for line in self.line_list]
        }
        return sum(self.chi2_dict.values())

Class variables

var channels : int
var chemical_params : dict
var chemistry_class : Type[ChemistryBase]

Subclass of ChemistryWB2014 that has non-zero default midplane and atmosphere abundance

Usage:

>>> physics = WilliamsBest2014(radial_bins=3, vertical_bins=3)
>>> chemistry = NonzeroChemistryWB2014(physics)
>>> chemistry.run_chemistry()
>>> chemistry.table["CO"].info.format = "e"
>>> chemistry.table["CO"]
<Column name='CO' dtype='float64' format='e' length=9>
1.000000e-10
1.000000e-10
1.000000e-10
1.000000e-04
1.000000e-10
1.000000e-10
1.000000e-10
1.000000e-10
1.000000e-10
var directory : Union[str, pathlib.Path]
var disk : str
var distance : Unit("pc")
var dust_opacity_file : Union[str, pathlib.Path]
var inclination : Unit("deg")
var line_list : List[Line]
var line_window_width : Unit("km / s")
var mctherm_threads
var mstar : Unit("solMass")
var nphot_therm : int
var npix : int
var physics_class : Type[PhysicsModel]

A subclass of WilliamsBest2014 which uses temperatures at 100 au rather than at 1 au for a more robust fitting

These defintions are equivalent:

>>> t_a, t_m, sl = 1000, 200, 0.55
>>> original = WilliamsBest2014(vertical_bins=3, radial_bins=3,
...     midplane_temperature_1au=t_m * u.K, atmosphere_temperature_1au=t_a * u.K,
...     temperature_slope=sl)
>>> at_100_au = WilliamsBest100au(vertical_bins=3, radial_bins=3,
...     midplane_temperature_100au=t_m * u.K / 100**sl, atmosphere_temperature_100au = t_a * u.K / 100**sl,
...     temperature_slope=sl)
>>> original.table     # doctest: +NORMALIZE_WHITESPACE
<CTable length=9>
   Radius       Height    Height to radius Gas density  Dust density Gas temperature Dust temperature
     AU           AU                         g / cm3      g / cm3           K               K
  float64      float64        float64        float64      float64        float64         float64
------------ ------------ ---------------- ------------ ------------ --------------- ----------------
1.000000e-01 0.000000e+00     0.000000e+00 1.153290e-15 1.153290e-17    7.096268e+02     7.096268e+02
1.000000e-01 3.500000e-02     3.500000e-01 5.024587e-34 5.024587e-36    3.548134e+03     3.548134e+03
1.000000e-01 7.000000e-02     7.000000e-01 5.921268e-72 5.921268e-74    3.548134e+03     3.548134e+03
7.071068e+00 0.000000e+00     0.000000e+00 2.285168e-13 2.285168e-15    6.820453e+01     6.820453e+01
7.071068e+00 2.474874e+00     3.500000e-01 2.716386e-17 2.716386e-19    3.410227e+02     3.410227e+02
7.071068e+00 4.949747e+00     7.000000e-01 7.189026e-23 7.189026e-25    3.410227e+02     3.410227e+02
5.000000e+02 0.000000e+00     0.000000e+00 2.871710e-20 2.871710e-22    6.555359e+00     6.555359e+00
5.000000e+02 1.750000e+02     3.500000e-01 6.929036e-22 6.929036e-24    2.625083e+01     2.625083e+01
5.000000e+02 3.500000e+02     7.000000e-01 8.170464e-23 8.170464e-25    3.277680e+01     3.277680e+01
>>> str(at_100_au.table) == str(original.table)
True
var physics_params : dict
var position_angle : Unit("deg")
var radial_bins_rt : int
var radmc_lines_run_kwargs : dict
var rstar : Unit("AU")
var tstar : Unit("K")
var vertical_bins_rt : int

Instance variables

var dust_directory
Expand source code
@property
def dust_directory(self):
    return self.directory / "radmc_dust"
var gas_directory
Expand source code
@property
def gas_directory(self):
    return self.directory / "radmc_gas"

Methods

def add_co_isotopologs(self)
Expand source code
def add_co_isotopologs(self, _13c: float = 77, _18o: float = 560):
    self.disk_chemical_model.table['13CO'] = self.disk_chemical_model.table['CO'] / _13c
    self.disk_chemical_model.table['C18O'] = self.disk_chemical_model.table['CO'] / _18o
    self.disk_chemical_model.table['13C18O'] = self.disk_chemical_model.table['CO'] / (_13c * _18o)
    # remember to fix the next abundance
    # self.disk_chemical_model.table['C17O'] = self.disk_chemical_model.table['CO'] / (
    #         560 * 5)  # http://articles.adsabs.harvard.edu/pdf/1994ARA%26A..32..191W
def chi2(self, uvs: Dict[str, UVFits])

Args

uvs
dictionary in a form of {line.name: uvt for line is self.line_list}

Returns

sum of chi2 between uvs and lines

Expand source code
def chi2(self, uvs: Dict[str, UVFits]):
    """

    Args:
        uvs: dictionary in a form of {line.name: uvt for line is self.line_list}

    Returns:
        sum of chi2 between uvs and lines
    """

    self.chi2_dict = {
        name: uv.chi2_with(self.gas_directory / f"{name}_image.fits")
        for name, uv in uvs.items()
        if name in [line.name for line in self.line_list]
    }
    return sum(self.chi2_dict.values())
def initialize_chemistry(self)
Expand source code
def initialize_chemistry(self):
    self.disk_chemical_model = self.chemistry_class(physics=self.disk_physical_model, **self.chemical_params)
def initialize_dust(self)
Expand source code
def initialize_dust(self):
    self.dust = DustPopulation(self.dust_opacity_file,
                               table=self.disk_physical_model.table,
                               name="Dust")
    self.dust.write_to_table()
def initialize_physics(self)
Expand source code
def initialize_physics(self):
    self.disk_physical_model = self.physics_class(**self.physics_params)
def mctherm(self, threads=None)

Run thermal radiative transfer for dust temperature calculation

Expand source code
def mctherm(self, threads=None):
    """Run thermal radiative transfer for dust temperature calculation"""
    if threads is None:
        threads = self.mctherm_threads
    folder_rt_dust = self.dust_directory
    folder_rt_dust.mkdir(parents=True, exist_ok=True)

    self.mctherm_model = RadMCTherm(
        chemistry=self.disk_chemical_model,
        folder=folder_rt_dust,
        nphot_therm=self.nphot_therm,
        star_radius=self.rstar,
        star_effective_temperature=self.tstar
    )

    self.mctherm_model.create_files()
    self.mctherm_model.run(threads=threads)
    self.mctherm_model.read_dust_temperature()
def run_chemistry(self)
Expand source code
def run_chemistry(self):
    self.disk_chemical_model.run_chemistry()
    self.add_co_isotopologs()
def run_line_radiative_transfer(self, run_kwargs: dict = None, **kwargs)

Run line radiative transfer

Expand source code
def run_line_radiative_transfer(
        self,
        run_kwargs: dict = None,
        **kwargs
):
    """Run line radiative transfer"""
    folder_rt_gas = self.gas_directory
    folder_rt_gas.mkdir(parents=True, exist_ok=True)

    disk_map = RadMCRTLines(
        chemistry=self.disk_chemical_model, line_list=self.line_list,
        radii_bins=self.radial_bins_rt, theta_bins=self.vertical_bins_rt,
        folder=folder_rt_gas, **kwargs
    )

    disk_map.create_files(channels_per_line=self.channels, window_width=self.line_window_width)

    disk_map.run(
        inclination=self.inclination,
        position_angle=self.position_angle,
        distance=self.distance,
        npix=self.npix,
        **self.radmc_lines_run_kwargs
    )