G = 1.
AU = 1.496e13 class Grid(object):
    """Construct a simple 1D grid with different spacings"""
    def __init__(self, R0, R1, N, spacing='log'):

        if spacing == 'log':
            self._setup_log(R0, R1, N)
        elif spacing == 'linear':
            self._setup_powerlaw(R0, R1, N, 1.0)
        elif spacing == 'natural':
            self._setup_powerlaw(R0, R1, N, 0.5)
        else:
            try:
                self._setup_powerlaw(R0, R1, N, float(spacing))
            except ValueError:
                raise AttributeError("Spacing must be a power law index, or "
                                     "one of 'log', 'linear' and 'natural'")
        self._setup_aux()
        self._N = N

        self._R0 = R0
        self._R1 = R1
        self._spacing = spacing
            
    def _setup_log(self, R0, R1, N):
        """Setup a grid in log-spacing"""
        lgR0 = np.log10(R0)
        lgR1 = np.log10(R1)
        
        dlogR = (lgR1 - lgR0) / N
        
        Ree = np.power(10, lgR0 + np.arange(-2, N+3, dtype='f8') * dlogR)
        
        self._Re  = Ree[2:-2]
        self._Rce = np.sqrt(Ree[2:-1] * Ree[1:-2])
        self._Rc  = self._Rce[1:-1]

        self._Ree = Ree

    def _setup_powerlaw(self, R0, R1, N, alpha):
        """Setup a power law grid"""
        alpha = float(alpha)
        alpha1 = 1/alpha

        R0a = R0**alpha
        R1a = R1**alpha

        dRa = (R1a - R0a) / N

        Ree_a = R0a + np.arange(-2, N+3, dtype='f8') * dRa
        Rce_a = 0.5*(Ree_a[2:-1] + Ree_a[1:-2])

        Ree = Ree_a**alpha1
        Rce = Rce_a**alpha1

        self._Re  = Ree[2:-2]
        self._Rce = Rce
        self._Rc  = Rce[1:-1]

        self._Ree = Ree
        
    def _setup_aux(self):
        self._dRe  = np.diff(self._Re)
        self._dRc  = np.diff(self._Rc)
        self._dRce = np.diff(self._Rce)

        self._dRe2  = np.diff(self._Re**2)
        self._dRc2  = np.diff(self._Rc**2)
        self._dRce2 = np.diff(self._Rce**2)

    def ASCII_header(self):
        """Write grid info header"""
        head = '# {} R0: {}, R1: {}, N: {}, spacing: {}'
        return head.format(self.__class__.__name__,
                           self._R0, self._R1, self.Ncells, self._spacing)

    def HDF5_attributes(self):
        """Class information for HDF5 headers"""
        def fmt(x):  return "{}".format(x)
        return self.__class__.__name__, { "R0" : fmt(self._R0),
                                          "R1" : fmt(self._R1),
                                          "N"  : fmt(self.Ncells),
                                          "spacing" : fmt(self._spacing),
                                          }
        
    @property
    def Rc(self):
        return self._Rc
    @property
    def Re(self):
        return self._Re
    @property
    def Rce(self):
        return self._Rce
    @property
    def Ree(self):
        return self._Ree

    @property
    def dRc(self):
        return self._dRc
    @property
    def dRe(self):
        return self._dRe
    @property
    def dRce(self):
        return self._dRce

    @property
    def dRc2(self):
        return self._dRc2
    @property
    def dRe2(self):
        return self._dRe2
    @property
    def dRce2(self):
        return self._dRce2

    @property
    def Ncells(self):
        return self._N

    def interp_centre(self, R, data):
        """Interpolate data to new radii

        args:
            R    : new radii
            data : data defined at grid centres
        """
        return np.interp(R, self.Rc, data)


    def interp_edges(self, R, data):
        """Interpolate data to new radii

        args:
            R    : new radii
            data : data defined at grid edges
        """
        return np.interp(R, self.Re, data)


    @staticmethod
    def from_string(string):
        """Read a Grid from a string"""
        string = string.replace('# Grid', '').strip()
        kwargs = {}
        args = [None, None, None]
        for item in string.split(','):
            key, val = [ x.strip() for x in item.split(':')]

            if   key == 'R0':
                args[0] = float(val)
            elif key == 'R1':
                args[1] = float(val)
            elif key == 'N':
                args[2] = int(val)
            elif key == 'spacing':
                kwargs[key] = val
            else:
                raise AttributeError("Error: Attribute {} for Grid not "
                                     "known".format(key))
        return Grid(*args, **kwargs) class SimpleStar(StarBase):
    """Simple non-evolving star."""
    def __init__(self, **kwargs):
        super(SimpleStar, self).__init__(**kwargs)

    def evolve(self, age, M=None):
        """Update the stellar properties based on current age and mass

        args:
           age : stellar age in yr
           M   : mass, Msun
        """
        self._age = age
        if M is not None: self._M = M

    @staticmethod
    def from_string(string):
        """Read a simple star from a string"""
        string = string.replace('# SimpleStar', '').strip()
        kwargs = {}
        for item in string.split(','):
            key, val = [ x.strip() for x in item.split(':')]

            if   key == 'M':
                kwargs[key] = float(val.replace('Msun','').strip())
            elif key == 'R':
                kwargs[key] = float(val.replace('Rsun','').strip())
            elif key == 'T' or key == 'T_eff':
                kwargs['T_eff'] = float(val.replace('K','').strip())
            elif key == 'age':
                kwargs[key] = float(val.replace('yr', '').strip())
            else:
                raise AttributeError("Error: Attribute {} for SimpleStar not "
                                     "known".format(key))
        return SimpleStar(**kwargs) class PhotoStar(SimpleStar):
    def __init__(self, LX=1e30, Phi=0, **kwargs):
        super().__init__(**kwargs)
        self._L_X = LX
        self._Phi = Phi

    @property
    def L_X(self):
        """X-ray Luminosity"""
        return self._L_X

    @property
    def Phi(self):
        """EUV Photon Luminosity"""
        return self._Phi class IrradiatedEOS(EOS_Table):
    """Model for an active irradiated disc.

    From Nakamoto & Nakagawa (1994), Hueso & Guillot (2005).

    args:
        star    : Stellar properties
        alpha_t : Viscous alpha parameter
        Tc      : External irradiation temperature (nebular), default=10
        Tmax    : Maximum temperature allowed in the disc, default=1500
        mu      : Mean molecular weight, default = 2.4
        gamma   : Ratio of specific heats
        kappa   : Opacity, default=Zhu2012
        accrete : Whether to include heating due to accretion,
                  default=True
    """
    def __init__(self, star, alpha_t, Tc=10, Tmax=1500., mu=2.4, gamma=1.4,
                 kappa=None,
                 accrete=True, tol=None): # tol is no longer used
        super(IrradiatedEOS, self).__init__()

        self._star = star
        
        self._dlogHdlogRm1 = 2/7.

        self._alpha_t = alpha_t
        
        self._Tc = Tc
        self._Tmax = Tmax
        self._mu = mu

        self._accrete = accrete
        self._gamma = gamma

        if kappa is None:
            self._kappa = opacity.Zhu2012
        else:
            self._kappa = kappa
        
        self._T = None

        self._compute_constants()

    def _compute_constants(self):
        self._sigTc4 = sig_SB*self._Tc**4
        self._cs0 = (Omega0**-1/AU) * (GasConst / self._mu)**0.5
        self._H0  = (Omega0**-1/AU) * (GasConst / (self._mu*self._star.M))**0.5


    def update(self, dt, Sigma, amax=1e-5, star=None):
        if star:
            self._star = star
            self._compute_constants()
        star = self._star
            
        # Temperature/gensity independent quantities:
        R = self._R
        Om_k = Omega0 * star.Omega_k(R)

        X = star.Rau/R
        f_flat  = (2/(3*np.pi)) * X**3
        f_flare = 0.5 * self._dlogHdlogRm1 * X**2
        
        # Heat capacity
        mu = self._mu
        #C_V = (k_B / (self._gamma - 1)) * (1 / (mu * m_H))
        
        alpha = self._alpha_t
        if not self._accrete:
            alpha = 0.

        # Local references 
        max_heat = sig_SB * (self._Tmax*self._Tmax)*(self._Tmax*self._Tmax)
        star_heat = sig_SB * star.T_eff**4
        sqrt2pi = np.sqrt(2*np.pi)            
        def balance(Tm):
            """Thermal balance"""
            cs = np.sqrt(GasConst * Tm / mu)
            H = cs / Om_k

            kappa = self._kappa(Sigma / (sqrt2pi * H), Tm, amax)
            tau = 0.5 * Sigma * kappa
            H /= AU

            # External irradiation
            dEdt = self._sigTc4
            
            # Compute the heating from stellar irradiation
            dEdt += star_heat * (f_flat + f_flare * (H/R))

            # Viscous Heating
            visc_heat = 1.125*alpha*cs*cs * Om_k
            dEdt += visc_heat*(0.375*tau*Sigma + 1./kappa)
            
            # Prevent heating above the temperature cap:
            dEdt = np.minimum(dEdt, max_heat)

            # Cooling
            Tm2 = Tm*Tm
            dEdt -= sig_SB * Tm2*Tm2

            # Change in temperature
            return (dEdt/Omega0) # / (C_V*Sigma)

        # Solve the balance using brent's method (needs ~ 20 iterations)
        T0 = self._Tc
        T1 = self._Tmax
        if self._T is not None:
            dedt = balance(self._T)
            T0 = np.where(dedt > 0, self._T, T0)
            T1 = np.where(dedt < 0, self._T, T1)

        self._T =  brentq(balance, T0, T1)
        self._Sigma = Sigma

        # Save the opacity:
        cs = np.sqrt(GasConst * self._T / mu)
        H = cs / Om_k
        self._kappa_arr = self._kappa(Sigma / (sqrt2pi * H), self._T, amax)
        
        self._set_arrays()


    def set_grid(self, grid):
        self._R = grid.Rc
        self._T = None

    def _set_arrays(self):
        super(IrradiatedEOS,self)._set_arrays()
        self._Pr = self._f_Pr()
    
    def __H(self, R, T):
        return self._H0 * np.sqrt(T * R*R*R)

    def _f_cs(self, R):
        return self._cs0 * self._T**0.5

    def _f_H(self, R):
        return self.__H(R, self._T)
    
    def _f_nu(self, R):
        return self._alpha_t * self._f_cs(R) * self._f_H(R)

    def _f_alpha(self, R):
        return self._alpha_t

    def _f_Pr(self):
        kappa = self._kappa_arr
        tau = 0.5 * self._Sigma * kappa
        f_esc = 1 + 2/(3*tau*tau)
        Pr_1 =  2.25 * self._gamma * (self._gamma - 1) * f_esc
        return 1. / Pr_1

    @property
    def T(self):
        return self._T

    @property
    def Pr(self):
        return self._Pr

    @property
    def star(self):
        return self._star

    def ASCII_header(self):
        """IrradiatedEOS header"""
        head = super(IrradiatedEOS, self).ASCII_header()
        head += ', opacity: {}, T_extern: {}K, accrete: {}, alpha: {}'
        head += ', Tmax: {}K'
        return head.format(self._kappa.__class__.__name__,
                           self._Tc, self._accrete, self._alpha_t,
                           self._Tmax)

    def HDF5_attributes(self):
        """Class information for HDF5 headers"""
        name, head = super(IrradiatedEOS, self).HDF5_attributes()

        head["opacity"]  = self._kappa.__class__.__name__
        head["T_extern"] = "{} K".format(self._Tc)
        head["accrete"]  = "{}".format(bool(self._accrete))
        head["alpha"]    = "{}".format(self._alpha_t)
        head["Tmax"]     = "{} K".format(self._Tmax)

        return name, head

    @staticmethod
    def from_file(filename):
        import star

        star = star.from_file(filename)
        alpha = None

        with open(filename) as f:
            for line in f:
                if not line.startswith('#'):
                    raise AttributeError("Error: EOS type not found in header")
                elif "IrradiatedEOS" in line:
                    string = line 
                    break 
                else:
                    continue

        kwargs = {}
        for item in string.split(','):    
            key, val = [ x.strip() for x in item.split(':')]

            if   key == 'gamma' or key == 'mu':
                kwargs[key] = float(val.strip())
            elif key == 'alpha':
                alpha = float(val.strip())
            elif key == 'accrete':
                kwargs[key] = bool(val.strip())
            elif key == 'T_extern':
                kwargs['Tc'] = float(val.replace('K','').strip())

        return IrradiatedEOS(star, alpha, **kwargs) class LocallyIsothermalEOS(EOS_Table):
    """Simple locally isothermal power law equation of state:

    args:
        h0      : aspect ratio at 1AU
        q       : power-law index of sound-speed
        alpha_t : turbulent alpha parameter
        star    : stellar properties
        mu      : mean molecular weight, default=2.4
    """
    def __init__(self, star, h0, q, alpha_t, mu=2.4):
        super(LocallyIsothermalEOS, self).__init__()
        
        self._h0 = h0
        self._cs0 = h0 * star.M**0.5
        self._q = q
        self._alpha_t = alpha_t
        self._H0 = h0
        self._T0 = (AU*Omega0)**2 * mu / GasConst
        self._mu = mu
        
    def _f_cs(self, R):
        return self._cs0 * R**self._q

    def _f_H(self, R):
        return self._H0 * R**(1.5+self._q)
    
    def _f_nu(self, R):
        return self._alpha_t * self._f_cs(R) * self._f_H(R)

    def _f_alpha(self, R):
        return self._alpha_t

    @property
    def T(self):
        return self._T0 * self.cs**2

    @property
    def Pr(self):
        return np.zeros_like(self._R)

    def ASCII_header(self):
        """LocallyIsothermalEOS header string"""
        head = super(LocallyIsothermalEOS, self).ASCII_header()
        head += ', h0: {}, q: {}, alpha: {}'
        return head.format(self._h0, self._q, self._alpha_t)

    def HDF5_attributes(self):
        """Class information for HDF5 headers"""
        name, head = super(LocallyIsothermalEOS, self).HDF5_attributes()
        head["h0"]   = "{}".format(self._h0)
        head["q"]     = "{}".format(self._q)
        head["alpha"] = "{}".format(self._alpha_t)
        return name, head

    @staticmethod
    def from_file(filename):
        raise NotImplementedError('')

    @property
    def star(self):
        return self._star class DustGrowthTwoPop(DustyDisc):
    """Two-population dust growth model of Birnstiel (2011).

    This model computes the flux of two dust populations. The smallest size
    particles are assumed to always be well coupled to the gas. For the larger
    particles we solve their growth up to the most stringent limit set by
    radial drift and fragmentation.

    Any dust tracers are assumed to have the same mass distribution as the dust
    particles themselves.

    args:
        grid      : Disc gridding object
        star      : Stellar object
        eos       : Equation of state
        eps       : Initital dust fraction
        Sigma     : Initial surface density distribution
        rho_s     : solid density, default=1
        Sc        : Schmidt number, default=1
        rhos      : Grain solid density, default=1.
        uf_0      : Fragmentation velocity (default = 100 (cm/s))
        uf_ice    : Fragmentation velocity of icy grains (default = 1000 (cm/s))
        f_ice     : Ice fraction, default=1
        thresh    : Threshold ice fraction for switchng between icy/non icy
                    fragmentation velocity, default=0.1
        f_grow    : Growth time-scale factor, default=1.
        a0        : Initial particle size (default = 1e-5, 0.1 micron)
        amin      : Minimum particle size (default = 0.0)
        f_drift   : Drift fitting factor. Reduce by a factor ~10 to model the
                    role of bouncing (default=0.55).
        f_frag    : Fragmentation boundary fitting factor (default=0.37).
        feedback  : Whether to include feedback from dust on gas
        start_small:Whether to start at monomer size (True, default) or equilibrium (False)
        distribution_slope:
                    The slope d ln n(a) / d ln a of the number distribution with size (3.5 for MRN)
    """
    def __init__(self, grid, star, eos, eps, Sigma=None,
                 rho_s=1., Sc=1., uf_0=100., uf_ice=1e3, f_ice=1, thresh=0.1,
                 f_grow=1.0, a0=1e-5, amin=1e-5, f_drift=0.55, f_frag=0.37, feedback=True, start_small=True, distribution_slope=3.5):
        super(DustGrowthTwoPop, self).__init__(grid, star, eos, Sigma, rho_s, Sc, feedback)
        
        self._uf_0   = uf_0 / (AU * Omega0)
        self._uf_ice = uf_ice / (AU * Omega0)

        # Fitting factors
        self._fgrow  = f_grow 
        self._ffrag  = f_frag * (2/(3*np.pi)) 
        self._fdrift = f_drift * (2/np.pi) / f_grow 
        self._fmass  = np.array([0.97, 0.75])

        # Initialize the dust distribution
        Ncells = self.Ncells
        self._fm    = np.zeros(Ncells, dtype='f8')
        self._a0    = 0 # Force well-coupled limit
        self._eps   = np.empty([2, Ncells], dtype='f8')
        self._a     = np.empty([2, Ncells], dtype='f8')
        self._eps[0] = eps
        self._eps[1] = 0
        self._a[0]   = amin
        self._a[1]   = a0
        
        self._amin = amin 
        
        self._ice_threshold = thresh
        self._uf = self._frag_velocity(f_ice)
        self._area = np.pi * a0*a0
        self._start_small = start_small         # Whether to start at monomer size (True, default) or equilibrium (False)
        self._p = distribution_slope            # The slope d ln n(a) / d ln a of the number distribution with size (3.5 for MRN)

        self._head = (', uf_0: {}cm s^-1, uf_ice: {}cm s^-1, thresh: {}'
                      ', f_grow: {}, a0: {}cm'.format(uf_0, uf_ice, thresh,
                                                      f_grow, a0))

        self.update(0)

    def ASCII_header(self):
        """Dust growth header"""
        return super(DustGrowthTwoPop, self).ASCII_header() + self._head

    def HDF5_attributes(self):
        """Class information for HDF5 headers"""
        name, head = super(DustGrowthTwoPop, self).HDF5_attributes()

        tmp = dict([x.strip().split(":") for x in self._head.split(",") if x])

        head.update(tmp)

        return self.__class__.__name__, head

    def _frag_velocity(self, f_ice):
        """Fragmentation velocity"""
        # Interplate between the icy/ice free region
        f_ice = np.minimum(f_ice/self._ice_threshold, 1)
        f_ice = f_ice*f_ice*f_ice*(10-f_ice*(15-6*f_ice))
        #f_ice = f_ice*f_ice*(3-2*f_ice)
        return self._uf_0 + (self._uf_ice - self._uf_0) * f_ice
        
    def _frag_limit(self):
        """Maximum particle size before fragmentation kicks in"""
        alpha = self.alpha/self.Sc
        af = (self.Sigma_G/(self._rho_s*alpha)) * (self._uf/self.cs)**2
        return self._ffrag * af

    def a_BT(self, eps_tot=None):
        """Size at transition between Brownian motion and turbulence dominated
        collision velocities"""
        if eps_tot is None:
            eps_tot = self.dust_frac.sum(0)

        alpha = self.alpha/self.Sc

        a0  = 8 * self.Sigma / (np.pi * self._rho_s) * self.Re**-0.25
        a0 *= np.sqrt(self.mu*m_H/(self._rho_s*alpha)) / (2*np.pi)
        return a0**0.4
        
    def _gammaP(self):
        """Dimensionless pressure gradient"""
        P = self.P
        R = self.R
        gamma = np.empty_like(P)
        gamma[1:-1] = abs((P[2:] - P[:-2])/(R[2:] - R[:-2]))
        gamma[ 0]   = abs((P[ 1] - P[  0])/(R[ 1] - R[ 0]))
        gamma[-1]   = abs((P[-1] - P[ -2])/(R[-1] - R[-2]))
        gamma *= R/(P+1e-300)

        return gamma
        
    def _drift_limit(self, eps_tot):
        """Maximum size due to drift limit or drift driven fragmentation"""
        gamma = self._gammaP()
        
        Sigma_D = self.Sigma * eps_tot
        Sigma_G = self.Sigma_G
            
        # Radial drift time-scale limit
        h = self.H / self.R
        ad = self._fdrift * (Sigma_D/self._rho_s) / (gamma * h**2+1e-300)

        # Radial drift-driven fragmentation:
        cs = self.cs
        St_d = 2 * (self._uf/cs) / (gamma*h + 1e-300)
        af = St_d * (2/np.pi) * (Sigma_G / self._rho_s)

        return ad, af

    def _t_grow(self, eps):
        return self._fgrow / (self.Omega_k * eps)

    def do_grain_growth(self, dt):
        """Apply the grain growth"""

        # Size and total gas fraction
        a = self._a[1]        
        eps_tot = self.dust_frac.sum(0)
                
        afrag_t = self._frag_limit()
        adrift, afrag_d =  self._drift_limit(eps_tot)
        t_grow = self._t_grow(eps_tot)
        
        afrag = np.minimum(afrag_t, afrag_d)
        a0    = np.minimum(afrag, adrift)       # a0 is the lower of the maximum sizes

        # Update the particle distribution
        #   Maximum size due to growth:
        if self._start_small:
            amax = np.minimum(a0, a*np.exp(dt/t_grow))  # If dust grains start small (default) first have to grow)
        else:
            amax = a0                                   # Ignore possibility of being in growth phase
        #   Reduce size due to erosion / fragmentation if grains have grown
        #   above this due to ice condensation
        # amin = a + np.minimum(0, afrag-a)*np.expm1(-dt/t_grow)
        # ignore empty cells:
        ids = eps_tot > 0
        self._a[1, ids] = np.maximum(amax[ids], self._amin)
        
        # Update the mass-fractions in each population
        fm   = self._fmass[1*(afrag < adrift)]
        self._fm[ids] = fm[ids]
        
        self._eps[0][ids] = ((1-fm)*eps_tot)[ids]
        self._eps[1][ids] = (   fm *eps_tot)[ids]

        # Set the average area:
        #self._area = np.pi * self.a_BT(eps_tot)**2

    def update_ices(self, grains):
        """Update the grain size due to a change in bulk ice abundance"""
        eps_new = grains.total_abund
            
        #f = eps_new / (self.integ_dust_frac + 1e-300)
        #self._a[1] = np.maximum(self._a0, self._a[1]*f**(1/3.))

        self._eps[0] = eps_new*(1-self._fm)
        self._eps[1] = eps_new*   self._fm

        # Update the ice fraction
        f_ice = 0
        for spec in grains:
            if 'grain' not in spec:
                f_ice += grains[spec]
        f_ice /= (eps_new + 1e-300)

        self._uf = self._frag_velocity(f_ice)

    def initialize_dust_density(self, dust_frac):
        """Set the initial dust density"""
        self._eps[0] = dust_frac


    def update(self, dt):
        """Do the standard disc update, and apply grain growth"""
        super(DustGrowthTwoPop, self).update(dt)
        self.do_grain_growth(dt) class Tazzari2016(object):
    '''Tabulated opacities for icy grains.

    The grain composition is taken from Tazzari et al. (2016, [1]).

    args:
        dust_to_gas : float, default = 0.01
            Dust-to-gas mass ratio
        q : float, default = 3.
            Slope of the grain size distribution, n(a) da ~ a^-q da

    References:
        [1] Tazzari, M., Testi, L., Ercolano, B., et al. 2016, A&A, 588, A53.
    '''
    def __init__(self, dust_to_gas = 0.01, q=3.5):
        data_dir = os.path.join(os.path.dirname(__file__), 'data', 'opacity')


        kappa_table_abs = 'kappa_table_abs_q{}.npy'.format(q)
        kappa_table_sca = 'kappa_table_sca_q{}.npy'.format(q)
        
        a         = np.load(os.path.join(data_dir, 'amax.npy'))
        T         = np.load(os.path.join(data_dir, 'T.npy'))
        kappa_abs = np.load(os.path.join(data_dir, kappa_table_abs)).T
        kappa_sca = np.load(os.path.join(data_dir, kappa_table_sca)).T

        self._amin = a[0]
        self._amax = a[-1]
        self._Tmin = T[0]
        self._Tmax = T[-1]

        self._kappa = RegularGridInterpolator((np.log(a), np.log(T)),
                                              np.log(kappa_abs + kappa_sca).T)

        self._eps = dust_to_gas

        self._q = q

        # Setup the coefficients for extrapolation
        k0_30, kl_30 = 3.84368330e+00, 1.15351466e-02
        k0_35, kl_35 = 3.94169075e+00, 2.27306292e-02

        self._k0 = k0_30 + (k0_35-k0_30) * ((q-3)/0.5)
        self._kl = kl_30 + (kl_35-kl_35) * ((q-3)/0.5)

    def __call__(self, rho, T, a):
        """Evaluate the opacity

        args:
            rho : float or array of float, ignored
                Density, c.g.s
            T : float or array of float
                Temperature, K
            amax : float, or array of float
                Maximum grain size, cm       
        """
        # Clip the limits to the tabulated ranges
        #  T < Tmin should never be needed, and kappa(T) ~ const at high T.
        #  a < amin kappa ~ const. and a > amax, we scale

        a_clip = np.minimum(np.maximum(a, self._amin), self._amax)
        T_clip = np.minimum(np.maximum(T, self._Tmin), self._Tmax)

        kappa = np.exp(self._kappa((np.log(a_clip), np.log(T_clip))))

        # Scale for large grains
        amax = np.maximum(a, a_clip)
        kp = self._q - self._k0 + self._kl*np.log10(amax/100)
        kappa *= (amax/a_clip)**kp

        return kappa * (self._eps / 0.01) class ViscousEvolution(object):
    """Solves the 1D viscous evolution equation.

    This class handles the inclusion of dust species in the one-fluid
    approximation. The total surface density (Sigma = Sigma_G + Sigma_D) is
    updated under the action of viscous forces, which are taken to act only
    on the gas phase species.

    Can optionally update tracer species.

    args:
       tol       : Ratio of the time-step to the maximum stable one.
                   Default = 0.5
       in_bound  : Type of internal boundary condition:
                     'Zero'      : Zero torque boundary
                     'Mdot'      : Constant Mdot (power law).
       boundary  : Type of external boundary condition:
                     'Zero'      : Zero torque boundary
                     'power_law' : Power-law extrapolation
                     'Mdot_inn'  : Constant Mdot, same as inner (power law).
                     'Mdot_out'  : Constant Mdot, for tail of LBP.
    """

    def __init__(self, tol=0.5, boundary='power_law', in_bound='Mdot'):
        self._tol = tol
        self._bound = boundary
        self._in_bound = in_bound

    def ASCII_header(self):
        """header"""
        return '# {} tol: {}'.format(self.__class__.__name__, self._tol)

    def HDF5_attributes(self):
        """Class information for HDF5 headers"""
        return self.__class__.__name__, { "tol" : str(self._tol)}



    def _setup_grid(self, grid):
        """Compute the grid factors"""
        self._X = 2 * np.sqrt(grid.Rc)
        self._dXe = 2 * np.diff(np.sqrt(grid.Re))
        self._dXc = 2 * np.diff(np.sqrt(grid.Rce))
        self._RXdXe = grid.Rc * self._X * self._dXe

    def _init_fluxes(self, disc):
        """Cache the important variables needed for the fluxes"""
        nuX = disc.nu * self._X

        S = np.zeros(len(nuX) + 2, dtype='f8')
        S[1:-1] = disc.Sigma_G * nuX

        # Inner boundary
        if self._in_bound == 'Zero':            # Zero torque
            S[0] = 0
        elif self._in_bound == 'Mdot':          # Constant flux (appropriate for power law)
            S[0] = S[1] * self._X[0] / self._X[1]
        else:
            raise ValueError("Error boundary type not recognised")

        # Outer boundary
        if self._bound == 'Zero':               # Zero torque
            S[-1] = 0
        elif self._bound == 'power_law':
            S[-1] = S[-2] ** 2 / S[-3]
        elif self._bound == 'Mdot_out':         # Constant flux (appropriate for tail of LBP)
            S[-1] = S[-2] * self._X[-2] / self._X[-1]
        elif self._bound == 'Mdot_inn':         # Constant flux (appropriate for power law)
            S[-1] = S[-2] * self._X[-1] / self._X[-2]
        else:
            raise ValueError("Error boundary type not recognised")

        self._dS = np.diff(S) / self._dXc

    def _fluxes(self):
        """Compute the mass fluxes for the viscous evolution equations.

        Gas update from Bath & Pringle (1981)
        """
        return 3. * np.diff(self._dS) / self._RXdXe

    def _tracer_fluxes(self, tracers):
        """Compute fluxes of a tracer.

        Uses the viscous update to compute the flux of  Sigma*tracers,
        divide by the updated Sigma to get the new tracer value.
        """
        shape = tracers.shape[:-1] + (tracers.shape[-1] + 2,)
        s = np.zeros(shape, dtype='f8')
        s[..., 1:-1] = tracers
        s[..., 0] = s[..., 1]
        s[..., -1] = s[..., -2]

        # Upwind the tracer density 
        ds = self._dS * np.where(self._dS <= 0, s[..., :-1], s[..., 1:])

        # Compute the viscous update
        return 3. * np.diff(ds) / self._RXdXe

    def viscous_velocity(self, disc, Sigma=None):
        """Compute the radial velocity due to viscosity"""
        self._setup_grid(disc.grid)
        self._init_fluxes(disc)

        # Can accept other density specifications (e.g. gas only to match specification of dust drift by Dipierro+18)
        if Sigma is None:
            Sigma = disc.Sigma
               
        RS = Sigma * disc.R
        return np.nan_to_num(- 3 * self._dS[1:-1] / (RS[1:] + RS[:-1])) # Avoid nans when working with Sigma_G -> 0

    def max_timestep(self, disc):
        """Courant limited time-step"""
        grid = disc.grid
        nu = disc.nu

        dXe2 = np.diff(2 * np.sqrt(grid.Re)) ** 2

        tc = ((dXe2 * grid.Rc) / (2 * 3 * nu)).min()
        return self._tol * tc

    def __call__(self, dt, disc, tracers=[]):
        """Compute one step of the viscous evolution equation
        args:
            dt      : time-step
            disc    : disc we are updating
            tracers : Tracer species to update. Should be a list of arrays with
                      shape = [*, disc.Ncells].
        """
        self._setup_grid(disc.grid)
        self._init_fluxes(disc)

        f = self._fluxes()
        Sigma_new = disc.Sigma + dt * f
        
        for t in tracers:
            if t is None: continue
            tracer_density = t*disc.Sigma
            t[:] = (dt*self._tracer_fluxes(t) + tracer_density) / (Sigma_new + 1e-300)

        disc.Sigma[:] = Sigma_new class ViscousEvolutionFV(object):
    """Solves the 1D viscous evolution equation via a finite-volume method

    This class handles the inclusion of dust species in the one-fluid
    approximation. The total surface density (Sigma = Sigma_G + Sigma_D) is
    updated under the action of viscous forces, which are taken to act only
    on the gas phase species.

    Can optionally update tracer species.

    args:
       tol       : Ratio of the time-step to the maximum stable one.
                   Default = 0.5
       in_bound  : Type of internal boundary condition:
                     'Zero'      : Zero torque boundary
                     'Mdot'      : Constant Mdot (power law).
       boundary  : Type of external boundary condition:
                     'Zero'      : Zero torque boundary
                     'power_law' : Power-law extrapolation
                     'Mdot_inn'  : Constant Mdot, same as inner (power law).
                     'Mdot_out'  : Constant Mdot, for tail of LBP.
    """

    def __init__(self, tol=0.5, boundary='power_law', in_bound='Mdot'):
        self._tol = tol
        self._bound = boundary
        self._in_bound = in_bound

    def ASCII_header(self):
        """header"""
        return '# {} tol: {}'.format(self.__class__.__name__, self._tol)

    def HDF5_attributes(self):
        """Class information for HDF5 headers"""
        return self.__class__.__name__, { "tol" : str(self._tol)}

    def _setup_grid(self, grid):
        """Compute the grid factors"""
        self._Rh  = np.sqrt(grid.Rc)
        self._dR3 = np.diff(np.sqrt(grid.Rce)**3)

        self._dA = grid.Re
        self._dV = 0.5 * np.diff(grid.Re**2)

    def _init_fluxes(self, disc):
        """Cache the important variables needed for the fluxes"""
        nuRh = disc.nu * self._Rh

        S = np.zeros(len(nuRh) + 2, dtype='f8')
        S[1:-1] = disc.Sigma_G * nuRh

        # Inner boundary
        if self._in_bound == 'Zero':            # Zero torque
            S[0] = 0
        elif self._in_bound == 'Mdot':          # Constant flux (appropriate for power law)
            S[0] = S[1] * self._Rh[0] / self._Rh[1]
        else:
            raise ValueError("Error boundary type not recognised")

        # Outer boundary
        if self._bound == 'Zero':
            S[-1] = 0
        elif self._bound == 'power_law':
            S[-1] = S[-2] ** 2 / S[-3]
        elif self._bound == 'Mdot_out':         # Constant flux (appropriate for tail of LBP)
            S[-1] = S[-2] * self._Rh[-2] / self._Rh[-1]
        elif self._bound == 'Mdot_inn':         # Constant flux (appropriate for power law)
            S[-1] = S[-2] * self._Rh[-1] / self._Rh[-2]
        else:
            raise ValueError("Error boundary type not recognised")

        self._dS = 4.5 * np.diff(S) / self._dR3

    def _fluxes(self):
        """Compute the mass fluxes for the viscous evolution equations.

        Gas update from Bath & Pringle (1981)
        """
        return np.diff(self._dA * self._dS) / self._dV

    def _tracer_fluxes(self, tracers):
        """Compute fluxes of a tracer.

        Uses the viscous update to compute the flux of  Sigma*tracers,
        divide by the updated Sigma to get the new tracer value.
        """
        shape = tracers.shape[:-1] + (tracers.shape[-1] + 2,)
        s = np.zeros(shape, dtype='f8')
        s[..., 1:-1] = tracers
        s[..., 0] = s[..., 1]
        s[..., -1] = s[..., -2]

        # Upwind the tracer density 
        ds = self._dS * np.where(self._dS <= 0, s[..., :-1], s[..., 1:])

        # Compute the viscous update
        return np.diff(self._dA * ds) / self._dV

    def viscous_velocity(self, disc, S = None):
        """Compute the radial velocity due to viscosity"""
        self._setup_grid(disc.grid)
        self._init_fluxes(disc)

        if S is None:
            S = disc.Sigma 
        return - 0.5 * self._dS[1:-1] / (S[1:] + S[:-1])

    def max_timestep(self, disc):
        """Courant limited time-step"""
        grid = disc.grid
        nu = disc.nu

        dXe2 = np.diff(2 * np.sqrt(grid.Re)) ** 2

        tc = ((dXe2 * grid.Rc) / (2 * 3 * nu)).min()
        return self._tol * tc

    def __call__(self, dt, disc, tracers=[]):
        """Compute one step of the viscous evolution equation
        args:
            dt      : time-step
            disc    : disc we are updating
            tracers : Tracer species to update. Should be a list of arrays with
                      shape = [*, disc.Ncells].
        """
        self._setup_grid(disc.grid)
        self._init_fluxes(disc)

        f = self._fluxes()
        Sigma_new = disc.Sigma + dt * f

        for t in tracers:
            if t is None: continue
            t[:] += dt*(self._tracer_fluxes(t) - t*f) / (Sigma_new + 1e-300)

        disc.Sigma[:] = Sigma_new class AccretionDisc(object):

    def __init__(self, grid, star, eos, Sigma=None):
        self._grid = grid
        self._star = star
        self._eos  = eos
        self._FUV = 0.0
        if Sigma is None:
            Sigma = np.zeros_like(self.R)
        self._Sigma = Sigma
        
        """ Extra properties for dealing with half empty cells in timescale approach """
        self.mass_lost = 0.0
        self.tot_mass_lost = 0.0
        self.i_edge = -1

    def ASCII_header(self):
        """Write header information about the disc"""
        head  = ''
        head += self._grid.ASCII_header() + '\n'
        head += self._star.ASCII_header() + '\n'
        head += self._eos.ASCII_header()
        return head

    def HDF5_attributes(self):
        """Class information for HDF5 headers"""
        return self.__class__.__name__, dict([ self._grid.HDF5_attributes(),
                                               self._star.HDF5_attributes(),
                                               self._eos.HDF5_attributes() ])

    def set_FUV(self,FUV):
        """Update the external FUV flux irradiating the disc"""
        self._FUV = FUV

    @property
    def star(self):
        return self._star

    @property
    def R(self):
        """Cell centre radii"""
        return self._grid.Rc

    @property
    def R_edge(self):
        """Cell edge radii"""
        return self._grid.Re

    @property
    def grid(self):
        return self._grid

    @property
    def Sigma(self):
        """Surface density"""
        return self._Sigma

    @property
    def Sigma_G(self):
        """Gas surface density"""
        return self.Sigma 

    @property
    def cs(self):
        """Sound speed"""
        return self._eos.cs

    @property
    def T(self):
        """Temperature"""
        return self._eos.T

    @property
    def mu(self):
        return self._eos.mu
    
    @property
    def H(self):
        """Scale-height"""
        return self._eos.H

    @property
    def h(self):
        """Aspect ratio"""
        return self.H/self.R

    @property
    def H_edge(self):
        """Scale-height at cell edge"""
        return self._eos.H_edge

    @property
    def P(self):
        return self.midplane_gas_density * self.cs**2

    @property
    def midplane_gas_density(self):
        return self.Sigma_G / (np.sqrt(2*np.pi) * self.H * AU)

    @property
    def midplane_density(self):
        return self.Sigma / (np.sqrt(2*np.pi) * self.H * AU)

    @property
    def column_density(self):
        n = self.midplane_gas_density / (self._eos._mu * m_H)
        Re = self.R_edge * AU
        ndR = n * (Re[1:] - Re[:-1])
        N = np.cumsum(ndR)
        return N

    @property
    def Ncells(self):
        return self._grid.Ncells

    @property
    def alpha(self):
        return self._eos.alpha

    @property
    def nu(self):
        return self._eos.nu
    
    @property
    def Re(self):
        """Reynolds number"""
        return (self.alpha*self.Sigma_G*sig_H2) / (2*self._eos.mu*m_H)

    @property
    def Pr(self):
        """Prandtl number"""
        return self._eos.Pr

    @property
    def Omega_k(self):
        return self._star.Omega_k(self.R)

    @property
    def FUV(self):
        return self._FUV

    """Methods to determine global properties of a viscous accretion disc"""
    def Rout(self, thresh=1e-5):
        """Determine the outer radius via density threshold"""
        notempty = self.Sigma_G > thresh
        notempty_cells = self.R_edge[1:][notempty]
        if np.size(notempty_cells>0):
            R_outer = notempty_cells[-1]
        else:
            R_outer = 0.0

        return R_outer

    def RC(self):
        """Fit an LBP profile to the disc and return the scale radius"""
        not_empty = (self.R < self.Rout())
        popt,pcov = optimize.curve_fit(LBP_profile,self.R[not_empty],np.log(self.Sigma_G[not_empty]),p0=[100,0.01],maxfev=5000)
        return popt[0]

    def Mtot(self):
        """Determine the total mass"""
        Re = self.R_edge * AU
        dA = np.pi * (Re[1:] ** 2 - Re[:-1] ** 2)
        dM_tot = self.Sigma * dA
        M_tot = np.sum(dM_tot)
        return M_tot

    def Mdot(self,viscous_velocity):
        """Determine the viscous accretion rate"""
        M_visc_out = 2*np.pi * self.R[0] * self.Sigma[0] * viscous_velocity * (AU**2)
        Mdot = -M_visc_out*(yr/Msun)
        return Mdot

    """Other methods"""
    def set_surface_density(self, Sigma):
        self._Sigma[:] = Sigma

    def update(self, dt):
        """Update the disc properties and age"""

        new_age = self._star.age + dt/(2*np.pi)
        self._star.evolve(new_age)
        self._eos.update(dt, self.Sigma, self._star)

    def interp(self, R, data):
        """Interpolate disc data to new radii

        args:
            R    : new radii
            data : data defined at grid locations
        """
        return self.grid.interp_centre(R, data) class SingleFluidDrift(object):
    """Radial Drift in the single fluid approximation with the short friction
    time limit.

    This class computes the single-fluid update of the dust fraction,
        d(eps_i)/dt = - (1/Sigma) grad [Sigma eps_i (Delta v_i - eps Delta v)],
    which is a vertically integrated version of equation (98) of Laibe & Price
    (2014). Note that the time-derivative on the LHS is the Lagrangian
    derivative in centre of mass frame. If an Eulerian (fixed) grid is used the
    advection step must be handled seperately.

    The dust-gas relative velocity, Delta v_i, is calculated following
    Tanaka+ (2005).

    Note:
        This currently neglects the viscous velocity, which can be important
        for small grains.

    args:
        diffusion : Diffusion algorithm, default=None
        settling  : Include settling in the velocity calculation, default=False
        van_leer  : Use 2nd-order Van-Leer reconstruction, default=False
    """
    def __init__(self, diffusion=None, settling=False, van_leer=False):
        self._diffuse = diffusion
        self._settling = settling
        self._van_leer = van_leer

    def ASCII_header(self):
        """Radial drift header"""
        head = ''
        if self._diffuse:
            head += self._diffuse.ASCII_header() + '\n'
        head += ('# {} diffusion: {} settling: {} van-leer: {}'
                 ''.format(self.__class__.__name__,
                           self._diffuse is not None,
                           self._settling,
                           self._van_leer))
        return head

    def HDF5_attributes(self):
        """Class information for HDF5 headers"""
        head = { "diffusion" : "{}".format(self._diffuse is not None),
                 "settling"  : "{}".format(self._settling),
                 "van-leer"  : "{}".format(self._van_leer),
                 }
        if self._diffuse is not None:
            head.update(dict([self._diffuse.HDF5_attributes()]))
        return self.__class__.__name__ , head

    def max_timestep(self, disc, v_visc=None):
        step = np.inf
        Cou = 0.5       # Courant number
        
        dV = self._compute_deltaV(disc, v_visc)
        dVout = np.empty((dV.shape[0],dV.shape[-1]+2))
        dVout[:,1:-1] = dV
        dVout[:, 0] = dVout[:, 1]
        dVout[:,-1] = dVout[:,-2]
        dVtot = np.abs(dVout[:,1:]) + np.abs(dVout[:,:-1])  # Potentially a cell can lose dust in both directions, both should be included to ensure stability
        return Cou * (disc.grid.dRe / dVtot).min()
    
    def _donor_flux(self, Ree, deltaV_i, Sigma, eps_i):
        """Compute flux using Donor-Cell method"""
        # Add boundary cells        
        shape_v   = (eps_i.shape[-1]+1,)
        shape_rho = eps_i.shape[:-1] + (eps_i.shape[-1]+2,)
        
        dV_i = np.empty(shape_v, dtype='f8')
        dV_i[1:-1] = deltaV_i - self._epsDeltaV
        dV_i[ 0] = dV_i[ 1] 
        dV_i[-1] = dV_i[-2] 
            
        Sig_eps = np.zeros(shape_rho, dtype='f8')
        Sig_eps[...,1:-1] = Sigma*eps_i
            
        # Upwind the density
        dc = DonorCell(Ree[1:-1], 1)
        Sig_eps = dc(dV_i, Sig_eps)

        # Compute the fluxes
        flux = Sig_eps * dV_i
        return flux
        
    def _van_leer_flux(self, Ree, deltaV_i, Sigma, eps_i, dt):
        """Compute flux using Van-Leer reconstruction"""
        # Add boundary cells
        shape_v   = (eps_i.shape[-1]+3,)
        shape_rho = eps_i.shape[:-1] + (eps_i.shape[-1]+4,)
        
        dV_i = np.empty(shape_v, dtype='f8')
        dV_i[2:-2] = deltaV_i - self._epsDeltaV
        dV_i[  :2] = dV_i[2]
        dV_i[-2: ] = dV_i[-3] 
            
        Sig_eps = np.zeros(shape_rho, dtype='f8')
        Sig_eps[...,2:-2] = Sigma*eps_i

        Sig_eps[..., 1] = np.where(dV_i[ 1] < 0, Sig_eps[..., 2], 0)
        Sig_eps[...,-2] = np.where(dV_i[-2] > 0, Sig_eps[...,-3], 0)
        
        # Upwind the density
        vl = VanLeer(Ree, 1)
        Sig_eps = vl(dV_i, Sig_eps, dt)

        # Compute the fluxes
        flux = Sig_eps * dV_i[1:-1]
        return flux
        
    def _fluxes(self, disc, eps_i, deltaV_i, St_i, dt=0):
        """Update a quantity that moves with the gas/dust"""

        Sigma = disc.Sigma
        grid = disc.grid

        if self._van_leer:
            flux = self._van_leer_flux(grid.Ree, deltaV_i, Sigma, eps_i, dt)
        else:
            flux = self._donor_flux(grid.Ree, deltaV_i, Sigma, eps_i)
            
        # Do the update
        deps = - np.diff(flux*grid.Re) / ((Sigma+1e-300) * 0.5*grid.dRe2)

        if self._diffuse:
            St2 = St_i**2
            try:
                Sc = disc.Sc
            except ValueError:
                Sc = self._diffuse.Sc
            Sc = Sc * (0.5625/(1 + 4*St2) + 0.4375 + 0.25*St2)
            depsdiff = self._diffuse(disc, eps_i, Sc)
            deps += depsdiff

        return deps

    def _compute_deltaV(self, disc, v_visc=None):
        """Compute the total dust-gas background velocity"""

        Sigma  = disc.Sigma
        SigmaD = disc.Sigma_D
        SigmaG = disc.Sigma_G
        Om_k   = disc.Omega_k
        a      = disc.grain_size
        R      = disc.R

        # Average to cell edges:        
        Om_kav  = 0.5*(Om_k      [1:] + Om_k      [:-1])
        Sig_av  = 0.5*(Sigma     [1:] + Sigma     [:-1]) + 1e-300
        SigD_av = 0.5*(SigmaD[...,1:] + SigmaD[...,:-1])
        SigG_av = 0.5*(SigmaG[...,1:] + SigmaG[...,:-1])
        a_av    = 0.5*(a    [..., 1:] + a     [...,:-1])
        R_av    = 0.5*(R         [1:] + R         [:-1])

        #eps_geff = ((R*SigmaG)[...,1:] + (R*SigmaG)[...,:-1]) / ((R*Sigma)     [1:] + (R*Sigma)     [:-1] + 1e-300)

        # Compute the density factors needed for the effect of feedback on
        # the radial drift velocity.
        eps_av = 0.
        eps_g = 1.
        #   SigG_av = Sig_av
        if disc.feedback:
            # By default, use the surface density
            eps_av = SigD_av / Sig_av
            #eps_g = np.maximum(1 - eps_av.sum(0), 1e-300)
            eps_g = np.maximum(SigG_av / Sig_av, 1e-300)
            #eps_geff = np.maximum(eps_geff, 1e-300)
            #SigG_av = Sig_av * eps_g

            # Use the midplane density instead
            if self._settling:
                rhoD = disc.midplane_dust_density
                rhoG = disc.midplane_gas_density
                rhoD_av = 0.5 * (rhoD[...,1:] + rhoD[...,:-1])
                rhoG_av = 0.5 * (rhoG    [1:] + rhoG    [:-1])
                rho_av = rhoD_av.sum(0) + rhoG_av + 1e-300

                eps_av = rhoD_av / rho_av
                eps_g  = np.maximum(rhoG_av / rho_av, 1e-300)
                
        # Compute the Stokes number        
        St_av = disc.Stokes(SigG_av, a_av+1e-300)

        # Compute the lambda factors
        # DON'T Use lambda * eps_g instead of lambda to avoid 0/0 in D_1 when eps_g -> 0.
        la0, la1 = 0, 0 
        St_1 = 1 / (1 + St_av**2)
        if disc.feedback:
            la0 = (eps_av/eps_g / (1     + St_av** 2)).sum(0)
            la1 = (eps_av/eps_g / (St_av + St_av**-1)).sum(0)

        # Compute the gas velocities due to pressure (with feedback):
        rho = disc.midplane_gas_density
        dPdr = np.diff(disc.P) / disc.grid.dRc
        eta = - dPdr / (0.5*(rho[1:] + rho[:-1] + 1e-300)*Om_kav)

        D_1 = 1 / ((1 + la0)**2 + la1**2)
        u_gas =            la1  * eta * D_1
        v_gas = - 0.5*(1 + la0) * eta * D_1

        # Compute the gas velocities due to viscosity per Dipierro+18 (with feedback):
        if v_visc is not None:
            u_gas += (1 + la0) * v_visc * D_1 / eps_g
            v_gas += 0.5 * la1 * v_visc * D_1 / eps_g

        # Dust-gas relative velocities:
        DeltaV = (2*v_gas / (St_av + St_av**-1) 
                  - u_gas / (1     + St_av**-2))

        # epsDeltaV = v_COM - v_gas (= 0 if dust mass is neglected)
        if disc.feedback:
            self._epsDeltaV = (eps_av * DeltaV).sum(0)
        else:
            self._epsDeltaV = 0

        return DeltaV

    def __call__(self, dt, disc, gas_tracers=None, dust_tracers=None, v_visc=None):
        """Apply the update for radial drift over time-step dt"""
        eps = disc.dust_frac
        a   = disc.grain_size
        eps_inv = 1. / (eps.sum(0) + np.finfo(eps.dtype).tiny)

        # Compute the dust-gas relative velocity
        DeltaV = self._compute_deltaV(disc, v_visc=v_visc)
        
        # Compute and apply the fluxes
        if gas_tracers is not None:
            gas_tracers[:] += dt * self._fluxes(disc, gas_tracers, 0, 0, dt)

        # Update the dust fraction, size and tracers
        d_tr = 0
        for eps_k, dV_k, a_k, St_k in zip(disc.dust_frac, DeltaV,
                                          disc.grain_size, disc.Stokes()):

            if dust_tracers is not None:
                t_k = dust_tracers * eps_k * eps_inv
                d_tr  += dt*self._fluxes(disc, t_k, dV_k, St_k, dt)
                
            # multiply a_k by the dust-to-gas ratio, so that constant functions
            # are advected perfectly
            #eps_a = a_k * eps_k
            #eps_a +=  dt*self._fluxes(disc, eps_a, dV_k, St_k, dt)
            
            eps_k[:] += dt*self._fluxes(disc, eps_k, dV_k, St_k, dt)

            #a_k[:] = eps_a / (eps_k + 1e-300)

        if dust_tracers is not None:
            dust_tracers[:] += d_tr
            
    def radial_drift_velocity(self, disc):
        """Compute the radial drift velocity for the disc"""
        DeltaV = self._compute_deltaV(disc)
        return DeltaV - self._epsDeltaV class TracerDiffusion(object):
    """Diffusion of trace species in a turbulent background.
    args:
        Sc : Schmidt number, ratio of momentum diffusivity to mass diffusivity,
             default=1.
        limit : Whether to limit the diffusive velocity to the sound speed,
                default=False.
    """
    def __init__(self, Sc=1, limit=False):
        self._Sc = Sc
        self._limit = limit

    def ASCII_header(self):
        """Tracer diffusion header"""
        return '# {} Sc: {}, flux_limit: {}'.format(self.__class__.__name__,
                                                    self.Sc, self._limit)

    def HDF5_attributes(self):
        """Class information for HDF5 headers"""
        return (self.__class__.__name__,
                {"Sc" : "{}".format(self.Sc),
                 "flux_limit" : "{}".format(bool(self._limit))}
                )

    @staticmethod
    def _diffusive_flux(disc, eps_i, Sc):
        """Compute the diffusive flux for a tracer
        args:
            disc  : accretion disc object
            eps_i : Concentration of species that is diffusing
            Sc    : Schmidt number

        returns:
            flux_diffuse : diffusive flux at edges (between cells only)
        """
        D = disc.nu / Sc
        Sigma_G = disc.Sigma_G

        # Use geometric average to avoid problems at the edge of evaporating
        # regions., where Sigma_G = 0 (and eps_i is ill-defined)
        DSig = D*Sigma_G
        DSig = np.sqrt(np.maximum(DSig[1:]*DSig[:-1], 0))

        return - DSig * np.diff(eps_i) / disc.grid.dRc

    def _get_Schmidt(self, disc, Sc):
        if Sc is None:
            try:
                Sc = disc.Sc
            except AttributeError:
                Sc = self.Sc
        return Sc

    def max_timestep(self, disc, Sc=None):
        """Courant limited time-step"""
        grid = disc.grid
        D = disc.nu / self._get_Schmidt(disc, Sc)

        return (0.25 * np.diff(grid.Re)**2 / D).min()

    def __call__(self, disc, eps_i, Sc=None):
        """Compute the rate of change of the surface density due to diffusion.

        args:
            disc  : Disc object
            eps_i : Concentration of species that is diffusing
            Sc    : Schmidt number. If not provided the value provided by the
                    disc is tried, before falling back to self.Sc
        returns:
            dSigma_i/dt : Rate of change of density of tracer
        """
        Sc = self._get_Schmidt(disc, Sc)

        grid = disc.grid
        Sigma = disc.Sigma
        F = self._diffusive_flux(disc, eps_i, Sc)

        if self._limit:
            max_f = Sigma_G*eps_i*self._eos.cs+1e-300
            F /= 1 + abs(F)/(0.5*(max_f[1:] + max_f[:-1]))

        F *= grid.Re[1:-1]

        depsdt = np.zeros_like(eps_i)
        depsdt[...,1:]  += F / grid.dRe2[1:]
        depsdt[...,:-1] -= F / grid.dRe2[:-1]

        depsdt /= Sigma + 1e-300
        return depsdt

    @property
    def Sc(self):
        return self._Sc class DiscEvolutionDriver(object):
    """Driver class for full evolution model.

    Required Arguments:
        disc : Disc model to update

    Optional Physics update:
        dust             : Update the dust, i.e. radial drift
        gas              : Update due to gas effects, i.e. Viscous evolution
        diffusion        : Seperate diffusion update
        internal_photo   : Remove gas by internal photoevaporation
        photoevaporation : Remove gas by external photoevaporation
        chemistry        : Solver for the chemical evolution

    History:
        history          : Tracks values of key parameters over time

    Note: Diffusion is usually handled in the dust dynamics module

    Other options:
        t0  : Starting time, default = 0, code units
        t_out:Previous output times, default = None, years
    """

    def __init__(self, disc, gas=None, dust=None, diffusion=None, chemistry=None, ext_photoevaporation=None, int_photoevaporation=None, history=None, t0=0.):

        self._disc = disc

        self._gas       = gas
        self._dust      = dust
        self._diffusion = diffusion
        self._chemistry = chemistry
        self._external_photo = ext_photoevaporation
        self._internal_photo = int_photoevaporation

        self._history = history

        self._t = t0
        self._nstep = 0

    def __call__(self, tmax):
        """Evolve the disc for a single timestep

        args:
            dtmax : Upper limit to time-step

        returns:
            dt : Time step taken
        """
        disc = self._disc

        # Compute the maximum time-step
        dt = tmax - self.t
        if self._gas:
            dt = min(dt, self._gas.max_timestep(self._disc))
        if self._dust:
            v_visc  = self._gas.viscous_velocity(disc)
            dt = min(dt, self._dust.max_timestep(self._disc, v_visc))
            if self._dust._diffuse:
                dt = min(dt, self._dust._diffuse.max_timestep(self._disc))
        if self._diffusion:
            dt = min(dt, self._diffusion.max_timestep(self._disc))       
        if self._external_photo and hasattr(self._external_photo,"_density"): # If we are using density to calculate mass loss rates, we need to limit the time step based on photoevaporation
            (dM_dot, dM_gas) = self._external_photo.optically_thin_weighting(disc)
            Dt = dM_gas[(dM_dot>0)] / dM_dot[(dM_dot>0)]
            Dt_min = np.min(Dt)
            dt = min(dt,Dt_min)
        
		# Determine tracers for dust step
        gas_chem, ice_chem = None, None
        dust = None
        try:
            gas_chem = disc.chem.gas.data
            ice_chem = disc.chem.ice.data
        except AttributeError:
            pass

		# Do dust evolution        
        if self._dust:
            self._dust(dt, disc,
                       gas_tracers=gas_chem,
                       dust_tracers=ice_chem, v_visc=v_visc)

		# Determine tracers for gas steps
        try:
            gas_chem = disc.chem.gas.data
            ice_chem = disc.chem.ice.data
        except AttributeError:
            pass
        try:
            dust = disc.dust_frac
        except AttributeError:
            pass

        # Do Advection-diffusion update
        if self._gas:
            self._gas(dt, disc, [dust, gas_chem, ice_chem])

        if self._diffusion:
            if gas_chem is not None:
                gas_chem[:] += dt * self._diffusion(disc, gas_chem)
            if ice_chem is not None:
                ice_chem[:] += dt * self._diffusion(disc, ice_chem)
            if dust is not None:
                dust[:] += dt * self._diffusion(disc, dust)

        # Do external photoevaporation
        if self._external_photo:
            self._external_photo(disc, dt)

        # Do internal photoevaporation
        if self._internal_photo:
            self._internal_photo(disc, dt/yr, self._external_photo)

        # Pin the values to >= 0 and <=1:
        disc.Sigma[:] = np.maximum(disc.Sigma, 0)        
        try:
            disc.dust_frac[:] = np.maximum(disc.dust_frac, 0)
            disc.dust_frac[:] /= np.maximum(disc.dust_frac.sum(0), 1.0)
        except AttributeError:
            pass
        try:
            disc.chem.gas.data[:] = np.maximum(disc.chem.gas.data, 0)
            disc.chem.ice.data[:] = np.maximum(disc.chem.ice.data, 0)
        except AttributeError:
            pass

        # Chemistry
        if self._chemistry:
            rho = disc.midplane_gas_density
            eps = disc.dust_frac.sum(0)
            grain_size = disc.grain_size[-1]
            T = disc.T

            self._chemistry.update(dt, T, rho, eps, disc.chem, 
                                   grain_size=grain_size)

            # If we have dust, we should update it now the ice fraction has
            # changed
            disc.update_ices(disc.chem.ice)

        # Now we should update the auxillary properties, do grain growth etc
        disc.update(dt)

        self._t += dt
        self._nstep += 1
        return dt

    @property
    def disc(self):
        return self._disc

    @property
    def t(self):
        return self._t

    @property
    def num_steps(self):
        return self._nstep

    @property
    def gas(self):
        return self._gas
    @property
    def dust(self):
        return self._dust
    @property
    def diffusion(self):
        return self._diffusion
    @property
    def chemistry(self):
        return self._chemistry
    @property
    def photoevaporation_external(self):
        return self._external_photo
    @property
    def photoevaporation_internal(self):
        return self._internal_photo
    @property
    def history(self):
        return self._history

    def dump_ASCII(self, filename):
        """Write the current state to a file, including header information"""

        # Put together a header containing information about the physics
        # included
        head = ''
        if self._gas:
            head += self._gas.ASCII_header() + '\n'
        if self._dust:
            head += self._dust.ASCII_header() + '\n'
        if self._diffusion:
            head += self._diffusion.ASCII_header() + '\n'
        if self._chemistry:
            head += self._chemistry.ASCII_header() + '\n'
        if self._external_photo:
            head += self._external_photo.ASCII_header() + '\n'
        if self._internal_photo:
            head += self._internal_photo.ASCII_header() + '\n'

        # Write it all to disc
        io.dump_ASCII(filename, self._disc, self.t, head)

    def dump_hdf5(self, filename):
        """Write the current state in HDF5 format, with header information"""
        headers = []
        if self._gas:            headers.append(self._gas.HDF5_attributes())
        if self._dust:           headers.append(self._dust.HDF5_attributes())
        if self._diffusion:      headers.append(self._diffusion.HDF5_attributes())
        if self._chemistry:      headers.append(self._chemistry.HDF5_attributes())
        if self._external_photo: headers.append(self._external_photo.HDF5_attributes())
        if self._internal_photo: headers.append(self._internal_photo.HDF5_attributes())

        io.dump_hdf5(filename, self._disc, self.t, headers) class Event_Controller(object):
    """Handles booking keeping for events that occur at the specified times.

    Event types are specified through key word arguments, e.g.
        Event_Controller(save=output_times, plot=plot_times)
    where the values passed must be iterable lists of times.
    """
    def __init__(self, **events):
        
        self._events  = {}
        self._event_number = {}
        for key in events:
            self._events[key] = sorted(events[key])
            self._event_number[key] = 0

    def event_types(self):
        return self._events.keys()

    def event_times(self, event_type):
        """Return the times of the specified event type"""
        return self._events[event_type]

    def next_event_time(self, event_type=None):
        """Time of next event. 

        If no event type is specified, the next event of any time is returned
        """
        if event_type is not None:
            return self._next_event_time(event_type)
        else:
            # All events:
            t_next = np.inf
            for event in self._events:
                t_next = min(t_next, self._next_event_time(event))
            return t_next

    def _next_event_time(self, event):
        try:
            return self._events[event][0]
        except IndexError:
            return np.inf

    def next_event(self):
        """The type of the next event"""
        t_next = self.next_event_time()
        for event in self._events:
            if self._next_event_time(event) == t_next:
                return event    
        return None

    def check_event(self, t, event):
        """Has the next occurance of a specified event occured?"""
        try:
            return self._events[event][0] <= t
        except IndexError:
            return False

    def events_passed(self, t):
        """Returns a list of event types that have passed since last pop"""
        return [ e for e in self._events if self.check_event(t, e) ]

    def event_number(self, key):
        """The number of times the specified event has occurred"""
        return self._event_number[key]

    def pop_events(self, t, event_type=None):
        """Remove events that have passed.

        If no event type is specified, pop all event types
        """
        if event_type is not None:
            self._pop(t, event_type)
        else:
            for event in self._events:
                self._pop(t, event)

    def _pop(self, t, event):        
        try:
            while self.check_event(t, event):
                self._events[event].pop(0)
                self._event_number[event] += 1
        except IndexError:
            pass
        
    def finished(self):
        """Returns True all events have been popped"""
        for event in self._events:
            if len(self._events[event]) > 0:
                return False
        else:
            return True class DiscReader(Reader):
    """Read disc snaphshots from file"""
    def __init__(self, DIR, base='disc', type='hdf5'):
        
        if type.lower() == 'ascii':
            SnapType = AsciiDiscSnap
            extension =  '.dat'
        else:
            SnapType = Hdf5DiscSnap
            extension = '.h5'

        super(DiscReader, self).__init__(SnapType, DIR, base, extension) def mkdir_p(path):
    try:
        os.makedirs(path)
    except OSError as exc:  # Python >2.5
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise class EUVDiscAlexander(PhotoBase):
    def __init__(self, disc, Type='Primordial', R_hole=None):
        super().__init__(disc, Regime='EUV', Type=Type)
        
        # Parameters for mass loss profiles
        self._cs = 10                                           # Sound speed in km s^-1
        self._RG = disc.star.M / (self._cs*1e5 /Omega0/AU)**2   # Gravitational Radius in AU
        self._mu = 1.35
        self._aB = 2.6e-13                                      # Case B Recombination coeff. in cm^3 s^-1
        self._C1 = 0.14
        self._A  = 0.3423
        self._B  = -0.3612
        self._D  = 0.2457
        self._C2 = 0.235
        self._a  = 2.42

        h = disc.H/disc.R
        he = np.empty_like(disc.R_edge)
        he[1:-1] = 0.5*(h[1:] + h[-1:])
        he[0] = 1.5*h[0] - 0.5*h[1] 
        he[-1] = 1.5*h[-1] - 0.5*h[-2]
        self._h = he

        # If initiating with an Inner Hole disc, need to update properties
        if self._type == 'InnerHole':
            self._Hole = True
            self._R_hole = R_hole
            #self.get_Rhole(disc)

        # Run the mass loss rates to update the table
        self.Sigma_dot(disc.R_edge, disc.star)

    def mdot_XE(self, star, Mdot=0):
        # Store Mdot calculated from profile
        self._Mdot = Mdot  # In Msun/yr
        self._Mdot_true = self._Mdot

    def scaled_R(self, R, star):
        if self._type=='Primordial':
            return R / self._RG         # Normalise to RG
        elif self._type=='InnerHole':
            return R / self.R_inner()   # Normalise to inner edge
        else:
            return R                    # If unspecified, don't modify

    def R_inner(self):
        # Innermost mass loss
        if self._type=='Primordial':
            return 0.1 * self._RG   # Mass loss profile is only positive for >0.1 RG
        elif self._type=='InnerHole':
            return self._R_hole     # Mass loss profile applies outside hole
        else:
            return 0                # If unspecified, assume mass-loss possible throughout 

    def Sigma_dot_Primordial(self, R, star, ret=False):
        Sigmadot = np.zeros_like(R)
        x = self.scaled_R(R,star)
        where_photoevap = (x >= 0.1)    # No mass loss close to star

        # Equation A3
        nG = self._C1 * (3 * star.Phi / (4*np.pi * (self._RG*AU)**3 * self._aB))**(1/2)        # cm^-3
        # Equation A2
        n0 = nG * (2 / (x**7.5 + x**12.5))**(1/5)

        # Equation A4
        u1 = self._cs*1e5*yr/Omega0 * self._A * np.exp(self._B * (x-0.1)) * (x-0.1)**self._D  # cm yr^-1

        # Combine terms (Equation A1)
        Sigmadot[where_photoevap] = 2 * self._mu * m_H * (n0 * u1)[where_photoevap]  # g cm^-2 /yr
        Sigmadot = np.maximum(Sigmadot,0)

        # Work out total mass loss rate
        dMdot = 2*np.pi * R * Sigmadot
        Mdot  = np.trapz(dMdot,R)  # g yr^-1 (AU/cm)^2

        # Normalise, convert to cgs
        Mdot  = Mdot * AU**2/Msun   # g yr^-1

        # Store result
        self.mdot_XE(star, Mdot=Mdot)

        if ret:
            # Return unaveraged values at cell edges
            return Sigmadot
        else:
            # Store values as average of mass loss rate at cell edges
            self._Sigmadot = (Sigmadot[1:] + Sigmadot[:-1]) / 2


    def Sigma_dot_InnerHole(self, R, star, ret=False):
        Sigmadot = np.zeros_like(R)
        x = self.scaled_R(R,star)
        where_photoevap = (x > 1)    # No mass loss inside hole

        # Combine terms (Equation A5)
        Sigmadot[where_photoevap] = (2 * self._mu * m_H * self._C2 * self._cs*1e5*yr/Omega0 * (star.Phi / (4*np.pi * (self.R_inner()*AU)**3 * self._aB * self._h))**(1/2) * x**(-self._a))[where_photoevap]  # g cm^-2 /yr
        Sigmadot = np.maximum(Sigmadot,0)

        # Work out total mass loss rate
        dMdot = 2*np.pi * R * Sigmadot
        Mdot  = np.trapz(dMdot,R)  # g yr^-1 (AU/cm)^2

        # Normalise, convert to cgs
        Mdot  = Mdot * AU**2/Msun   # g yr^-1

        # Store result
        self.mdot_XE(star, Mdot=Mdot)

        # Mopping up in the gap
        mop_up = (R >= 0.1 * self._RG) * (x <= 1.0)
        Sigmadot[mop_up] = np.inf

        if ret:
            # Return unaveraged values at cell edges
            return Sigmadot
        else:
            # Store values as average of mass loss rate at cell edges
            self._Sigmadot = (Sigmadot[1:] + Sigmadot[:-1]) / 2 class XrayDiscOwen(PhotoBase):
    def __init__(self, disc, Type='Primordial', R_hole=None):
        super().__init__(disc, Regime='X-ray', Type=Type)

        # Parameters for Primordial mass loss profile
        self._a1 = 0.15138
        self._b1 = -1.2182
        self._c1 = 3.4046
        self._d1 = -3.5717
        self._e1 = -0.32762
        self._f1 = 3.6064
        self._g1 = -2.4918
        # Parameters for Inner Hole mass loss profile
        self._a2 = -0.438226
        self._b2 = -0.10658387
        self._c2 = 0.5699464
        self._d2 = 0.010732277
        self._e2 = -0.131809597
        self._f2 = -1.32285709

        # If initiating with an Inner Hole disc, need to update properties
        if self._type == 'InnerHole':
            self._Hole = True
            self._R_hole = R_hole
            #self.get_Rhole(disc)

        # Run the mass loss rates to update the table
        self.Sigma_dot(disc.R_edge, disc.star)

    def mdot_XE(self, star, Mdot=None):
        # In Msun/yr
        if Mdot is not None:
            self._Mdot = Mdot
        elif self._type=='Primordial':
            self._Mdot = 6.25e-9 * star.M**(-0.068) * (star.L_X / 1e30)**(1.14) # Equation B1
        elif self._type=='InnerHole':
            self._Mdot = 4.8e-9 * star.M**(-0.148) * (star.L_X / 1e30)**(1.14)  # Equation B4
        else:
            raise NotImplementedError("Disc is of unrecognised type, and no mass-loss rate has been manually specified")
        self._Mdot_true = self._Mdot

    def scaled_R(self, R, star):
        # Where R in AU
        x = 0.85 * R / star.M                       # Equation B3
        if self._Hole:
            y = 0.95 * (R-self._R_hole) / star.M    # Equation B6
        else:
            y = R
        return x, y

    def R_inner(self, star):
        # Innermost mass loss
        return 0.7 / 0.85 * star.M

    def Sigma_dot_Primordial(self, R, star, ret=False):
        # Equation B2
        Sigmadot = np.zeros_like(R)
        x, y = self.scaled_R(R,star)
        where_photoevap = (x >= 0.7) * (x<=99) # No mass loss close to star, mass loss prescription becomes negative at log10(x)=1.996
        logx = np.log(x[where_photoevap])
        log10 = np.log(10)
        log10x = logx/log10

        # First term
        exponent = self._a1 * log10x**6 + self._b1 * log10x**5 + self._c1 * log10x**4 + self._d1 * log10x**3 + self._e1 * log10x**2 + self._f1 * log10x + self._g1
        t1 = 10**exponent

        # Second term
        terms = 6*self._a1*logx**5/log10**7 + 5*self._b1*logx**4/log10**6 + 4*self._c1*logx**3/log10**5 + 3*self._d1*logx**2/log10**4 + 2*self._e1*logx/log10**3 + self._f1/log10**2
        t2 = terms/x[where_photoevap]**2

        # Third term
        t3 = np.exp(-(x[where_photoevap]/100)**10)

        # Combine terms
        Sigmadot[where_photoevap] = t1 * t2 * t3

        # Work out total mass loss rate for normalisation
        M_dot = 2*np.pi * R * Sigmadot
        total = np.trapz(M_dot,R)

        # Normalise, convert to cgs
        Sigmadot = np.maximum(Sigmadot,0)
        Sigmadot *= self.Mdot / total * Msun / AU**2 # in g cm^-2 / yr

        if ret:
            # Return unaveraged values at cell edges
            return Sigmadot
        else:
            # Store values as average of mass loss rate at cell edges
            self._Sigmadot = (Sigmadot[1:] + Sigmadot[:-1]) / 2

    def Sigma_dot_InnerHole(self, R, star, ret=False):
        # Equation B5
        Sigmadot = np.zeros_like(R)
        x, y = self.scaled_R(R,star)
        where_photoevap = (y >= 0.0) # No mass loss inside hole
        use_y = y[where_photoevap]

        # Exponent of second term
        exp2 = -(use_y/57)**10

        # Numerator
        terms = self._a2*self._b2 * np.exp(self._b2*use_y+exp2) + self._c2*self._d2 * np.exp(self._d2*use_y+exp2) + self._e2*self._f2 * np.exp(self._f2*use_y+exp2)

        # Divide by Denominator
        Sigmadot[where_photoevap] = terms/R[where_photoevap]

        # Work out total mass loss rate for normalisation
        M_dot = 2*np.pi * R * Sigmadot
        total = np.trapz(M_dot,R)

        # Normalise, convert to cgs
        Sigmadot = np.maximum(Sigmadot,0)
        Sigmadot *= self.Mdot / total * Msun / AU**2 # in g cm^-2 / yr

        # Mopping up in the gap
        mop_up = (x >= 0.7) * (y < 0.0)
        Sigmadot[mop_up] = np.inf

        if ret:
            # Return unaveraged values at cell edges
            return Sigmadot
        else:
            # Store values as average of mass loss rate at cell edges
            self._Sigmadot = (Sigmadot[1:] + Sigmadot[:-1]) / 2 class XrayDiscPicogna(PhotoBase):
    def __init__(self, disc, Type='Primordial', R_hole=None):
        super().__init__(disc, Regime='X-ray', Type=Type)

        # Parameters for Primordial mass loss profile
        self._a1 = -0.5885
        self._b1 = 4.3130
        self._c1 = -12.1214
        self._d1 = 16.3587
        self._e1 = -11.4721
        self._f1 = 5.7248
        self._g1 = -2.8562
        # Parameters for Inner Hole mass loss profile
        self._a2 = 0.11843
        self._b2 = 0.99695
        self._c2 = 0.48835

        # If initiating with an Inner Hole disc, need to update properties
        if self._type == 'InnerHole':
            self._Hole = True
            self._R_hole = R_hole
            #self.get_Rhole(disc)

        # Run the mass loss rates to update the table
        self.Sigma_dot(disc.R_edge, disc.star)

    def mdot_XE(self, star, Mdot=None):
        # In Msun/yr
        if Mdot is not None:
            self._Mdot = Mdot
        elif self._type=='Primordial':
            logMd = -2.7326 * np.exp((np.log(np.log(star.L_X)/np.log(10))-3.3307)**2/-2.9868e-3) - 7.2580  # Equation 5
            self._Mdot = 10**logMd
        elif self._type=='InnerHole':
            logMd = -2.7326 * np.exp((np.log(np.log(star.L_X)/np.log(10))-3.3307)**2/-2.9868e-3) - 7.2580  # 1.12 * Equation 5
            self._Mdot = 1.12 * (10**logMd)
        else:
            raise NotImplementedError("Disc is of unrecognised type, and no mass-loss rate has been manually specified")
        self._Mdot_true = self._Mdot

    def scaled_R(self, R, star):
        # Where R in AU
        # All are divided by stellar mass normalised to 0.7 Msun (value used by Picogna+19) to represent rescaling by gravitational radius 
        x = R / (star.M/0.7)
        if self._Hole:
            y = (R-self._R_hole) / (star.M/0.7)    # Equation B6
        else:
            y = R / (star.M/0.7)
        return x, y

    def R_inner(self, star):
        # Innermost mass loss
        if self._type=='Primordial':
            return 0                # Mass loss possible throughout
        elif self._type=='InnerHole':
            return self._R_hole     # Mass loss profile applies outside hole
        else:
            return 0                # If unspecified, assume mass loss possible throughout 

    def Sigma_dot_Primordial(self, R, star, ret=False):
        # Equation B2
        Sigmadot = np.zeros_like(R)
        x, y = self.scaled_R(R,star)
        where_photoevap = (x<=137) # Mass loss prescription becomes negative at x=1.3785
        logx = np.log(x[where_photoevap])
        log10 = np.log(10)
        log10x = logx/log10

        # First term
        exponent = self._a1 * log10x**6 + self._b1 * log10x**5 + self._c1 * log10x**4 + self._d1 * log10x**3 + self._e1 * log10x**2 + self._f1 * log10x + self._g1
        t1 = 10**exponent

        # Second term
        terms = 6*self._a1*log10x**5 + 5*self._b1*log10x**4 + 4*self._c1*log10x**3 + 3*self._d1*log10x**2 + 2*self._e1*log10x + self._f1
        t2 = terms/(2*np.pi*x[where_photoevap]**2)

        # Combine terms
        Sigmadot[where_photoevap] = t1 * t2

        # Work out total mass loss rate for normalisation
        M_dot = 2*np.pi * R * Sigmadot
        total = np.trapz(M_dot,R)

        # Normalise, convert to cgs
        Sigmadot = np.maximum(Sigmadot,0)
        Sigmadot *= self.Mdot / total * Msun / AU**2 # in g cm^-2 / yr

        if ret:
            # Return unaveraged values at cell edges
            return Sigmadot
        else:
            # Store values as average of mass loss rate at cell edges
            self._Sigmadot = (Sigmadot[1:] + Sigmadot[:-1]) / 2

    def Sigma_dot_InnerHole(self, R, star, ret=False):
        # Equation B5
        Sigmadot = np.zeros_like(R)
        x, y = self.scaled_R(R,star)
        where_photoevap = (y > 0.0) * (y < -self._c2/np.log(self._b2)) # No mass loss inside hole, becomes negative at x=-c/ln(b)
        use_y = y[where_photoevap]

        # Numerator
        terms = self._a2 * np.power(self._b2,use_y) * np.power(use_y,self._c2-1) * (use_y * np.log(self._b2) + self._c2)

        # Divide by Denominator
        Sigmadot[where_photoevap] = terms/(2*np.pi*R[where_photoevap])

        # Work out total mass loss rate for normalisation
        M_dot = 2*np.pi * R * Sigmadot
        total = np.trapz(M_dot,R)

        # Normalise, convert to cgs
        Sigmadot = np.maximum(Sigmadot,0)
        Sigmadot *= self.Mdot / total * Msun / AU**2 # in g cm^-2 / yr

        # Mopping up in the gap - assume usual primordial rates there.
        Sigmadot[(y<=0.0) * (x<=137)] = self.Sigma_dot_Primordial(R, star, ret=True)[(y<=0.0)*(x<=137)]/1.12 # divide by 1.12 so that normalise to correct mass loss rate
        mop_up = (x > 137) * (y < 0.0)
        Sigmadot[mop_up] = np.inf   # Avoid having discontinuous mass-loss by filling in the rest

        if ret:
            # Return unaveraged values at cell edges
            return Sigmadot
        else:
            # Store values as average of mass loss rate at cell edges
            self._Sigmadot = (Sigmadot[1:] + Sigmadot[:-1]) / 2 class History(object):

    """Setup"""
    def __init__(self, dust, dthresh):
        self._threshold = 1e-5          # Threshold for defining edge by density
        self._dust = dust
        if self._dust:
            self._dthresholds = dthresh # Threshold percentiles of dust mass
        
        self._times = np.array([])      # Output times

        """Radii"""
        self._Rout  = np.array([])      # Outer radius of the disc (density)
        self._Rc    = np.array([])      # Radius of current best fit scale radius
        self._Ropt  = np.array([])      # Radius where Mdot maximum ie where becomes optically thick
        self._Rh    = np.array([])      # Outer radius of the transition disc hole
        if dust:
            self._Rdust = {}            # Radius containing user defined fraction of dust
            for threshold in self._dthresholds:
                self._Rdust[threshold] = np.array([])

        """Mass"""
        self._Mtot = np.array([])       # Total mass
        if dust:
            self._Mdust = np.array([])  # Mass of dust in disc

        """Mass Loss"""
        self._Mdot_acc = np.array([])   # Accretion rate, updated with velocity passed
        self._Mdot_ext = np.array([])   # External photoevaporation rate
        self._Mdot_int = np.array([])   # Internal photoevaporation rate
        self._Mcum_gas = np.array([])   # Total amount of gas lost to wind
        if dust:
            self._Mcum_dust = np.array([])  # Total amount of dust lost to wind

    """Methods to return values"""
    # Return times
    @property
    def times(self):
        return self._times

    # Return radii/masses/mass loss rates
    @property
    def radii(self):
        return self._Rout, self._Rc, self._Ropt, self._Rh
    @property
    def radii_dust(self):
        if self._dust:
            return self._Rdust
        else:
            print("No dust radii to return")
            return None

    @property
    def mass(self):
        return self._Mtot, self._Mcum_gas
    @property
    def mass_dust(self):
        if self._dust:
            return self._Mdust, self._Mcum_dust
        else:
            print("No dust masses to return")
            return None

    @property
    def mdot(self):
        return self._Mdot_acc, self._Mdot_ext, self._Mdot_int
    
    """When restarting, for any variable found in input file, import data"""
    def restart(self, filename, snap_number):
        restartdata = np.genfromtxt(filename, names=True, comments='#', max_rows=snap_number+1)
        
        self._times = restartdata['t']
        print("Restarting with times:\n",self.times)
        # No exception - must have times!

        try:
            self._Rout = restartdata['R_out']
        except ValueError:
            pass
        try:
            self._Rc   = restartdata['R_C']
        except ValueError:
            pass
        try:
            self._Ropt = restartdata['R_out']
        except ValueError:
            pass
        try:
            self._Rh   = restartdata['R_hole']
        except ValueError:
            if 'M_int' in restartdata.dtype.names:
                self._Rh = np.full_like(self._times,np.nan)
            else:
                pass
        try:
            self._Mtot = restartdata['M_D']
        except ValueError:
            pass
        try:
            self._Mdot_acc = restartdata['M_acc']
        except ValueError:
            pass
        try:
            self._Mdot_ext = restartdata['M_ext']
        except ValueError:
            pass
        try:
            self._Mdot_int = restartdata['M_int']
        except ValueError:
            pass
        try:
            self._Mcum_gas = restartdata['M_wg']
        except ValueError:
            pass

        # Dust-Only Parameters
        if self._dust:
            for threshold in self._dthresholds:
                try:
                    self._Rdust[threshold] = restartdata['R_{}'.format(int(float(threshold)*100))]
                except ValueError:
                    pass
            try:
                self._Mdust = restartdata['M_d']
            except ValueError:
                pass
            try:
                self._Mcum_dust = restartdata['M_wd']
            except ValueError:
                pass

    """Call"""
    def __call__(self, driver):
        self._times = np.append(self._times,[driver.t / yr])

        # 1 Radius: Density threshold / Scale radius of LBP / Optically Thin
        self._Rout = np.append(self._Rout,[driver.disc.Rout(thresh=self._threshold)])
        if not (driver.photoevaporation_external or driver.photoevaporation_internal):
            self._Rc = np.append(self._Rc,[driver.disc.RC()])
        # Optically thin still needed here!

        # 2 Disc Mass: Total
        self._Mtot = np.append(self._Mtot,[driver.disc.Mtot()])

        # 3 Dust radii and mass
        if self._dust:
            dust_radii = driver.disc.Rdust(thresholds=self._dthresholds) # Radius containing proportion of dust mass
            for thresh in self._dthresholds:
                self._Rdust[thresh] = np.append(self._Rdust[thresh],dust_radii[thresh])
            self._Mdust = np.append(self._Mdust,[driver.disc.Mdust()])   # Remaining dust mass

        # 4 Viscous accretion rate
        if driver.gas:            
            self._Mdot_acc = np.append(self._Mdot_acc,[driver.disc.Mdot(driver._gas.viscous_velocity(driver.disc)[0])])

        # 5 External photoevaporation mass loss
        if driver.photoevaporation_external:
            self._Mdot_ext  = np.append(self._Mdot_ext,[driver.photoevaporation_external._Mdot])
            self._Ropt      = np.append(self._Ropt,[driver.photoevaporation_external._Rot])
            self._Mcum_gas  = np.append(self._Mcum_gas,[driver.photoevaporation_external._Mcum_gas])       # Total mass of gas  lost in wind
            if self._dust:
                self._Mcum_dust = np.append(self._Mcum_dust,[driver.photoevaporation_external._Mcum_dust]) # Total mass of dust lost in wind

        # 6 Internal photoevaporation
        if driver.photoevaporation_internal:
            R_hole, N_hole = driver.photoevaporation_internal.get_Rhole(driver.disc)
            self._Rh = np.append(self._Rh,[R_hole])
            self._Mdot_int = np.append(self._Mdot_int, [driver.photoevaporation_internal._Mdot_true])

    """Remove hole radii if identification of hole later deemed to be wrong"""
    def clear_hole(self):
        self._R = np.full_like(self._Rh, np.nan)

    """Save the history so far"""
    def save(self, driver, save_directory):
        # 0 Select times of recording
        used_times = self.times

        # 1 Retrieve radii
        outer_radii, scale_radii, opt_radii, hole_radii = self.radii
        radii_select = {}
        if driver.photoevaporation_external:
            radii_select['R_out'] = opt_radii
        else:
            radii_select['R_out'] = outer_radii
        if np.isnan(scale_radii).sum() < len(scale_radii):
            radii_select['R_C'] = scale_radii
        if np.isnan(hole_radii).sum() < len(hole_radii):
            radii_select['R_hole'] = hole_radii

        # 2 Retrieve masses
        disc_masses, gas_wind = self.mass

        # 3 Dust
        if driver.dust:
            dust_radii = self.radii_dust
            dust_masses, dust_wind = self.mass_dust

        # 4 Accretion rates
        Macc, Mext, Mint = self.mdot

        # 5 Photoevaporation rates
        Mevap = {}
        if driver.photoevaporation_external:
            Mevap['M_ext'] = Mext
        if driver.photoevaporation_internal:
            Mevap['M_int'] = Mint

        # Save data
        outputdata = np.column_stack((used_times, disc_masses))
        head  = ['t','M_D']
        units = ['yr','g']
        if driver.dust:
            outputdata = np.column_stack((outputdata, dust_masses))
            head.append('M_d')
            units.append('g')

        for key, radii in radii_select.items():
            outputdata = np.column_stack((outputdata, radii))
            head.append(key)
            units.append('AU')
        if driver.dust:
            for key, radii in dust_radii.items():
                outputdata = np.column_stack((outputdata, radii))
                head.append('R_{}'.format(int(float(key)*100)))
                units.append('AU')

        if driver.gas:
            outputdata = np.column_stack((outputdata, Macc))
            head.append('M_acc')
            units.append('Msun/yr')
        for key, mdot in Mevap.items():
            outputdata = np.column_stack((outputdata, mdot))
            head.append(key)
            units.append('Msun/yr')
        
        if driver.photoevaporation_external:
            outputdata = np.column_stack((outputdata, gas_wind))
            head.append('M_extg')
            units.append('g')
        if driver.photoevaporation_external and driver.dust:
            outputdata = np.column_stack((outputdata, dust_wind))
            head.append('M_extd')
            units.append('g')

        head  = "\t".join(head)
        units = "\t".join(units)
        full_head = "\n".join([head,units])
            
        np.savetxt(str(save_directory)+"/discproperties.dat", outputdata, delimiter='\t', header=full_head)

        return outputdata               
import os
import json
import numpy as np
import matplotlib.pyplot as plt
import DiscEvolution.photoevaporation as photoevaporation
import DiscEvolution.FRIED.photorate as photorate
    import argparse
from DiscEvolution.constants import Msun, AU, yr, Mjup
from DiscEvolution.grid import Grid
from DiscEvolution.star import SimpleStar, PhotoStar
from DiscEvolution.eos  import IrradiatedEOS, LocallyIsothermalEOS
from DiscEvolution.dust import DustGrowthTwoPop
from DiscEvolution.opacity import Tazzari2016
from DiscEvolution.viscous_evolution import ViscousEvolution, ViscousEvolutionFV
from DiscEvolution.disc import AccretionDisc
from DiscEvolution.dust import SingleFluidDrift
from DiscEvolution.diffusion import TracerDiffusion
from DiscEvolution.driver import DiscEvolutionDriver
from DiscEvolution.io import Event_Controller, DiscReader
from DiscEvolution.disc_utils import mkdir_p
from DiscEvolution.internal_photo import EUVDiscAlexander, XrayDiscOwen, XrayDiscPicogna
from DiscEvolution.history import History

# run_model.py
#
# Author: R. Booth
# Date: 4 - Jun - 2018
#
# Run a disc evolution model with transport and absorption / desorption but
# no other chemical reactions. 
#
# Note:
#   The environment variable 
#       "KROME_PATH=/home/rab200/WorkingCopies/krome_ilee/build"
#   should be set.
###############################################################################


###############################################################################
# Global Constants
###############################################################################

DefaultModel = "DiscConfig_default.json"

def LBP_profile(R,R_C,Sigma_C):
    # For profile fitting
    x = R/R_C
    return np.log(Sigma_C) - np.log(x)-x

###############################################################################
# Setup Functions
###############################################################################

def setup_disc(model):
    '''Create disc object from initial conditions'''

    # Setup the grid
    p = model['grid']
    grid = Grid(p['R0'], p['R1'], p['N'], spacing=p['spacing'])

    # Setup the star with photoionizing luminosity if provided and non-zero
    p = model['star']
    try:
        if model['x-ray']['L_X'] > 0:
            star = PhotoStar(LX=model['x-ray']['L_X'], Phi=0, M=p['mass'], R=p['radius'], T_eff=p['T_eff'])
        elif model['euv']['Phi'] > 0:
            star = PhotoStar(LX=0, Phi=model['euv']['Phi'], M=p['mass'], R=p['radius'], T_eff=p['T_eff'])
        else:
            star = SimpleStar(M=p['mass'], R=p['radius'], T_eff=p['T_eff'])
    except KeyError:

        star = SimpleStar(M=p['mass'], R=p['radius'], T_eff=p['T_eff'])
