models.py

import numpy as np
from scipy import optimize as sciop
try:
  from numba import njit
  def cdec(func):
    return njit(func)
except ModuleNotFoundError:
  def cdec(func):
    return func

#physical constants
_MSUN = 1.989*10**33
_G = 6.6743*10**-8
_c = 2.99792458*10**10
_kb = .3807*10**-16
_sigma = 5.6704*10**-5
_mh = 1.6726*10**-24
_sigmaT = 6.6524587158*10**-25
_a = 4*_sigma/_c


@cdec
def _Omega(r, M, sigma):
  return np.sqrt(_G*M/r**3 + 2.0*(sigma/r)**2)

@cdec
def _psolve(T, mu, P, rho):
  p = rho*_kb*T/mu + _a*T**4/3
  return p-P

@cdec
def _pcheck(T, gscale, rscale, ptot):
  pguess = T*gscale + rscale*T**4
  return ptot - pguess

class SS73:
  def __init__(self,
               opacityFunction = None, 	# optional opacity function (defaults to Kramers)
               omegaSigma = 0.0, 	# optional velocity dispersion to include in Omega
               X = 0.72,		# H mass fraction
               Z = 0.02): 		# metal mass fraction
    self.X = X
    self.Z = Z
    self.mu = _mh*4/(3+5*X-Z)
    self.kes = 0.2*(1 + X)
    self.sigma = omegaSigma

    if opacityFunction is None:
      #@cdec
      def _opacKramer(rho, T):
        kk = 4.0*10**25*(self.X+1.0)*(self.Z + 0.001)*rho*T**-3.5
        return self.kes + kk
      self.opac = _opacKramer
    else:
      self.opac = opacityFunction

  def _rhosolve(self, rho_g, tau, T, B):
    kappa = self.opac(rho_g, T)
    return tau - kappa*B**(1./3.)*rho_g**(2./3.)

  def _tsolve(self, Tg, A, B, omega, mu, rho_g):
    tau = Tg**4/A
    res = sciop.fsolve(self._rhosolve, x0=rho_g, args = (tau, Tg, B), full_output=1)
    rho = res[0]
    H = (B/rho)**(1./3.)
    P = rho*H**2*omega**2
    res = sciop.fsolve(_psolve, x0 = Tg, args=(mu, P, rho), full_output=1)
    return res[0] - Tg

  def _rho_t_solve(self, x, A, B, omega, mu):
    rho = x[0]
    T = x[1]

    kappa = self.opac(rho, T)
    tau = T**4/A
    rho_c = np.sqrt((tau/kappa)**3/B)

    H = (B/rho)**(1./3.)
    P = rho*H**2*omega**2
    res = sciop.fsolve(_psolve, x0 = T, args=(mu, P, rho), full_output=1)
    T_c = res[0][0]

    return np.array([rho-rho_c, T-T_c])

  def _solveSS(self, omega, mdot, alpha, mu, T_g = None, rho_g=None):
    A = 9*mdot*omega**2/(np.pi*64*_sigma)
    B = mdot/(6*np.pi*alpha*omega)

    if T_g is None:
      T_g = omega*(3/_a)**0.5*(B/(A*self.kes))**0.25
    if rho_g is None:
      rho_g = omega**6*B*(3/(_a*A*self.kes))**3

    resT = sciop.fsolve(self._rho_t_solve, x0=(rho_g, T_g), args=(A, B, omega, mu), full_output=1)
    T = resT[0][1]
    rho = resT[0][0]

    H = (B/rho)**(1./3.)
    P = rho*H**2*omega**2
    tau = T**4/A

    ierr = resT[2]
    if ierr !=1:
      resT = sciop.fsolve(self._tsolve, x0=T_g, args=(A, B, omega, mu, rho_g), full_output=1)
      T = resT[0][0]
      tierr = resT[2]
      tau = T**4/A
      resrho = sciop.fsolve(self._rhosolve, x0=rho_g, args = (tau, T_g, B), full_output=1 )
      rho = resrho[0][0]
      rhoierr = resrho[2]
      ierr = tierr or rhoierr

      H = (B/rho)**(1./3.)
      P = rho*H**2*omega**2

    return T, rho, H, P, tau, ierr

  def calcModel(self, logm, alpha, eta, eps = 0.1, rmin = 6.0, rin = 10.0, rout = 10**5, Nr = 10000):
    M = 10.0**logm
    MBH = M*_MSUN
    Rg = _G*MBH/_c**2
    Mdot0 = eta*4.0*np.pi*_G*MBH*_mh/(eps*_c*_sigmaT)
    rs = np.geomspace(rin, rout, Nr)*Rg
    Mdot = Mdot0*(1.0 - np.sqrt(rmin*Rg/rs) )
    Omega = _Omega(rs, MBH, self.sigma)

    ts = np.empty(Nr)
    rhos = np.empty(Nr)
    hs = np.empty(Nr)
    ps = np.empty(Nr)
    taus = np.empty(Nr)
      
    badTs = np.empty(Nr)
    badRhos = np.empty(Nr)

    for i in range(Nr):
      if i==0:
        Ti, rhoi, Hi, Pi, taui, terr = self._solveSS(Omega[i], Mdot[i], alpha, self.mu)
      else:
        if badTs[i-1]==1:
          tguess = ts[i-1]
          rguess = rhos[i-1]
        else:
          lastGoodT = np.nonzero(badTs==1)[0][-1]
          tguess = ts[lastGoodT]
          rguess = rhos[lastGoodT]
        Ti, rhoi, Hi, Pi, taui, terr = self._solveSS(Omega[i], Mdot[i], alpha, self.mu, tguess, rguess)

      hs[i] = Hi
      taus[i] = taui
      ps[i] = Pi
      badTs[i] = terr
      ts[i] = Ti
      rhos[i] = rhoi
      
    good = (badTs==1) #&(badRhos==1)

    Rs = rs[good]
    Ts = ts[good]
    Rhos = rhos[good]
    Hs = hs[good]
    Taus = taus[good]
    Ps = ps[good]
    Omegas = Omega[good]

    Ts = np.interp(rs, Rs, Ts)
    Rhos = np.interp(rs, Rs, Rhos)
    Hs = np.interp(rs, Rs, Hs)
    Taus = np.interp(rs, Rs, Taus)
    Ps = np.interp(rs, Rs, Ps)
    Omegas = np.interp(rs, Rs, Omegas)

    Kappas = Taus/(Rhos*Hs)
    Cs = Hs*Omegas
    Sigmas = 2.0*Hs*Rhos
    Qs = Cs*Omegas/(np.pi*_G*Sigmas)

    Pgas = Rhos*_kb*Ts/self.mu
    Prad = _a*Ts**4/3

    model = {"Rg":              Rg,
             "badfrac":         1-np.sum(good)/Nr,
             "logM":            logm, 

             "badcells":        1-good,
             "radius":          rs,
             "dens":            Rhos,
             "surfDens":        Sigmas,
             "temp":            Ts,
             "height":          Hs,
             "tau":             Taus,
             "press":           Ps,
             "prad":            Prads,
             "pgas":            Pgass,
             "omega":           Omegas,
             "kappa":           Kappas,
             "cs":              Cs,
             "toomreQ":         Qs }


    Pgass = Rhos*_kb*Ts/self.mu
    Prads = (_sigma*Taus*0.5/_c)*Ts**4/(Taus*0.375 + 0.5 + 0.25/Taus)

    model = {"Rg":              Rg,
             "badfrac":         1-np.sum(good)/Nr,
             "logM":            logm, 

             "badcells":        1-good,
             "radius":          rs,
             "dens":            Rhos,
             "surfDens":        Sigmas,
             "temp":            Ts,
             "height":          Hs,
             "tau":             Taus,
             "press":           Ps,
             "prad":            Prads,
             "pgas":            Pgass,
             "omega":           Omegas,
             "kappa":           Kappas,
             "cs":              Cs,
             "toomreQ":         Qs }
    return model



    return model


class SG03:
  def __init__(self,
               opacityFunction = None, 	# optional opacity function (defaults to Kramers)
               omegaSigma = 0.0, 	# optional velocity dispersion to include in Omega
               X = 0.72,		# H mass fraction
               Z = 0.02, 		# metal mass fraction
               Qc = 1.0):
    self.Qc = Qc
    self.X = X
    self.Z = Z
    self.mu = _mh*4/(3+5*X-Z)
    self.kes = 0.2*(1 + X)
    self.sigma = omegaSigma

    if opacityFunction is None:
      def _opacKramer(rho, T):
        kk = 4.0*10**25*(self.X+1.0)*(self.Z + 0.001)*rho*T**-3.5
        return self.kes + kk
      self.opac = _opacKramer
    else:
      self.opac = opacityFunction

  def _solve_tau_cs(self, x, A, B, omega, mu):
    tau = x[0]
    cs = x[1]
    tauf = tau*0.375 + 0.5 + 0.25/tau
    temp = (tauf*A)**0.25
    rho = B/cs**3
    tau_check = rho*(cs/omega)*self.opac(rho, temp)

    cs_check = ( _kb*temp/mu + tau*_sigma*0.5*A/(_c*rho) )**0.5

    return np.array([tau - tau_check, cs - cs_check])

  def _solve_4(self, x, A, B, omega, mu):
    T = x[0]
    tau = x[1]
    cs = x[2]
    rho = x[3]

    tau_c = self.opac(rho, T)*rho*cs/omega
    cs_c = np.sqrt(_kb*T/mu + tau*A*_sigma*0.5/(rho*_c) )
    rho_c = B/cs**3
    T_c = ((tau*0.375 + 0.5 + 0.25/tau)*A)**0.25

    return np.array([T-T_c, tau-tau_c, cs - cs_c, rho - rho_c])

  def _solveSGstable(self, omega, mdot, alpha, mu, rho_g = None, T_g=None):
    A = 3*mdot*omega**2/(np.pi*8*_sigma)
    B = mdot*omega**2/(6.0*alpha*np.pi)

    if T_g is None or rho_g is None:
      cs_g = _sigma*self.kes*A*0.5/(_c*omega)
      rho_g = B/cs_g**3
      tau_g = rho_g*cs_g*self.kes/omega
      T_g = ((tau_g*0.375 + 0.5 + 0.25/tau_g)*A)**0.25
    else:
      kappa = self.opac(rho_g, T_g)
      cs_g = (B/rho_g)**(1./3.)
      tau_g = kappa*rho_g*cs_g/omega

    res4 = sciop.fsolve(self._solve_4, x0=(T_g, tau_g, cs_g, rho_g), args=(A, B, omega, mu), full_output=1)
    temp = res4[0][0]
    tau = res4[0][1]
    cs = res4[0][2]
    rho = res4[0][3]
    tau_ierr = res4[2]

    if tau_ierr != 1: #try optically thin gas pressure solution?
      T_g = A**0.25
      tau_g = 1.0
      cs_g = np.sqrt(_kb*T_g/mu)
      rho_g = B/cs_g**3

      res4 = sciop.fsolve(self._solve_4, x0=(T_g, tau_g, cs_g, rho_g), args=(A, B, omega, mu), full_output=1)
      temp = res4[0][0]
      tau = res4[0][1]
      cs = res4[0][2]
      rho = res4[0][3]
      tau_ierr = res4[2]

      if tau_ierr != 1: # probably inner region, electron scattering opacity but non-negligible gas pressure? 
        taus = np.geomspace(10, 10**8, 110)
        residuals = []
        for tau_g in taus:
          T_g = ((tau_g*0.375 + 0.5 + 0.25/tau_g)*A)**0.25
          cs_g = (B/A)*(2*_c)/(tau_g*_sigma)
          rho_g = B/cs_g**3

          arg = np.array([T_g, tau_g, cs_g, rho_g])
          diff = self._solve_4(arg, A, B, omega, mu)
          residuals.append(np.sum((diff/arg)**2))
        best = np.argmin(residuals)

        tau_g = taus[best]
        T_g = ((tau_g*0.375 + 0.5 + 0.25/tau_g)*A)**0.25
        cs_g = (B/A)*(2*_c)/(tau_g*_sigma)
        rho_g = B/cs_g**3

        res4 = sciop.fsolve(self._solve_4, x0=(T_g, tau_g, cs_g, rho_g), args=(A, B, omega, mu), full_output=1)
        temp = res4[0][0]
        tau = res4[0][1]
        cs = res4[0][2]
        rho = res4[0][3]
        tau_ierr = res4[2]

    H = cs/omega
    P = cs*cs*rho

    return temp, rho, H, P, tau, tau_ierr

  def _solve_temp(self, T, rho, h, ptot, gscale, rscale):
    kappa = self.opac(rho, T)
    tau = rho*kappa*h
    tauf = 0.375*tau + 0.5 + 0.25/tau
    pcheck = gscale*T + rscale*tau*T**4/tauf
    a = (pcheck/(rscale*tauf))**0.25

    rest = sciop.fsolve(_pcheck, x0=a, args = (gscale, rscale*tau/tauf, ptot), full_output=1 )
    temp = rest[0][0]
    if rest[2] != 1:
      rest = sciop.fsolve(_pcheck, x0=(pcheck/gscale), args = (gscale, rscale, ptot), full_output=1 )
      temp = rest[0][0]
    return T - temp

  def _solveSGunstable(self, omega, mdot, alpha, mu, T_g):
    rho = omega**2/(2.0*np.pi*_G*self.Qc)
    h = (mdot/(rho*6*np.pi*alpha*omega))**(1./3.)
    cs = h*omega

    ptot = cs*cs*rho
    rscale = _sigma*0.5/_c
    gscale = rho*_kb/mu

    resT = sciop.fsolve(self._solve_temp, x0=(T_g), args=(rho, h, ptot, gscale, rscale), full_output=1)
    temp = resT[0][0]
    t_ierr = resT[2]


    kappa = self.opac(rho, temp)
    tau = rho*kappa*h

    return temp, rho, h, ptot, tau, t_ierr



  def calcModel(self, logm, alpha, eta, eps = 0.1, rmin = 6.0, rin = 10.0, rout = 10**5, Nr = 10000):
    M = 10.0**logm
    MBH = M*_MSUN
    Rg = _G*MBH/_c**2
    Mdot0 = eta*4.0*np.pi*_G*MBH*_mh/(eps*_c*_sigmaT)
    rs = np.geomspace(rin, rout, Nr)*Rg
    Mdot = Mdot0*(1.0 - np.sqrt(rmin*Rg/rs) )
    Omega = _Omega(rs, MBH, self.sigma)

    ts = np.empty(Nr)
    rhos = np.empty(Nr)
    hs = np.empty(Nr)
    ps = np.empty(Nr)
    taus = np.empty(Nr)
      
    badTs = np.empty(Nr)
    stable = True
    for i in range(Nr):
      if stable:
        if i==0 or np.sum(badTs)<1:
          Ti, rhoi, Hi, Pi, taui, terr = self._solveSGstable(Omega[i], Mdot[i], alpha, self.mu)
        else:
          if badTs[i-1]==1: 
            tguess = ts[i-1]
            rhoguess = rhos[i-1]
            Ti, rhoi, Hi, Pi, taui, terr = self._solveSGstable(Omega[i], Mdot[i], alpha, self.mu, rhoguess, tguess)
          else:
            try:
              lastGoodT = np.nonzero(badTs==1)[0][-1]
              tguess = ts[lastGoodT]
              rhoguess = rhos[lastGoodT] 
              Ti, rhoi, Hi, Pi, taui, terr = self._solveSGstable(Omega[i], Mdot[i], alpha, self.mu, rhoguess, tguess) 
            except IndexError:
              Ti, rhoi, Hi, Pi, taui, terr = self._solveSGstable(Omega[i], Mdot[i], alpha, self.mu) 
        Qi = Omega[i]**2*0.5/(np.pi*_G*rhoi)
        if (Qi < self.Qc):
          stable = False
          Ti, rhoi, Hi, Pi, taui, terr = self._solveSGunstable(Omega[i], Mdot[i], alpha, self.mu, Ti)
      else:
        if badTs[i-1]==1: 
          tguess = ts[i-1]
        else:
          lastGoodT = np.nonzero(badTs==1)[0][-1]
          tguess = ts[lastGoodT]
        Ti, rhoi, Hi, Pi, taui, terr = self._solveSGunstable(Omega[i], Mdot[i], alpha, self.mu, tguess)

      hs[i] = Hi
      taus[i] = taui
      ps[i] = Pi
      badTs[i] = terr
      ts[i] = Ti
      rhos[i] = rhoi
 
    good = (badTs==1)

    Rs = rs[good]
    Ts = ts[good]
    Rhos = rhos[good]
    Hs = hs[good]
    Taus = taus[good]
    Ps = ps[good]
    Omegas = Omega[good]

    Ts = np.interp(rs, Rs, Ts)
    Rhos = np.interp(rs, Rs, Rhos)
    Hs = np.interp(rs, Rs, Hs)
    Taus = np.interp(rs, Rs, Taus)
    Ps = np.interp(rs, Rs, Ps)
    Omegas = np.interp(rs, Rs, Omegas)

    #after cleaning or not, calculate auxilliary quantities 

    Kappas = Taus/(Rhos*Hs)
    Cs = Hs*Omegas
    Sigmas = 2.0*Hs*Rhos
    Qs = Omegas**2/(2.0*np.pi*_G*Rhos)

    Pgass = Rhos*_kb*Ts/self.mu
    Prads = (_sigma*Taus*0.5/_c)*Ts**4/(Taus*0.375 + 0.5 + 0.25/Taus)

    model = {"Rg":              Rg,
             "badfrac":         1-np.sum(good)/Nr,
             "logM":            logm, 

             "badcells":        1-good,
             "radius":          rs,
             "dens":            Rhos,
             "surfDens":        Sigmas,
             "temp":            Ts,
             "height":          Hs,
             "tau":             Taus,
             "press":           Ps,
             "prad":            Prads,
             "pgas":            Pgass,
             "omega":           Omegas,
             "kappa":           Kappas,
             "cs":              Cs,
             "toomreQ":         Qs }
    return model