Example of auto_diff in run_star_extras

Example of auto_diff in run_star_extras#

The auto_diff module simplifies using run_star_extras hooks that need to provide partial derivatives. For example, the other_surface_PT hook can be used to specify an alternate surface boundary condition by providing the surface temperature and pressure. This hook has the following call signature:

other_surface_PT(id, &
        skip_partials, &
        lnT_surf, dlnT_dL, dlnT_dlnR, dlnT_dlnM, dlnT_dlnkap, &
        lnP_surf, dlnP_dL, dlnP_dlnR, dlnP_dlnM, dlnP_dlnkap, ierr)

In addition to lnT_surf and lnP_surf, we also need to provide four partial derivatives! We could specify these manually. For instance an Eddington atmosphere might have:

subroutine other_surface_PT(id, &
           skip_partials, &
           lnT_surf, dlnT_dL, dlnT_dlnR, dlnT_dlnM, dlnT_dlnkap, &
           lnP_surf, dlnP_dL, dlnP_dlnR, dlnP_dlnM, dlnP_dlnkap, ierr)

use const_def, only: dp
integer, intent(in) :: id
logical, intent(in) :: skip_partials
real(dp), intent(out) :: &
   lnT_surf, dlnT_dL, dlnT_dlnR, dlnT_dlnM, dlnT_dlnkap, &
   lnP_surf, dlnP_dL, dlnP_dlnR, dlnP_dlnM, dlnP_dlnkap
integer, intent(out) :: ierr

real(dp) :: g

lnT_surf = log(pow(s%L(1) / (4d0 * pi * boltzsig * pow2(s%R(1))), 0.25d0))
dlnT_dL = lnT_surf / s%L(1)
dlnT_dlnR = -0.5d0
dlnT_dlnM = 0d0
dlnT_dlnkap = 0d0

g = s%M(1) / pow2(s%R(1))
lnP_surf = log(g / s%opacity(1))
dlnP_dL = 0d0
dlnP_dlnR = -2d0
dlnP_dlnM = 1d0
dlnP_dlnkap = -1d0

end subroutine other_surface_PT

This relies on us calculating the derivatives correctly by hand and hard-coding them. For this simple case that isn’t too hard, but it is easy to imagine more complicated cases where doing it by hand is error-prone and tedious. Instead, we could use auto_diff:

subroutine other_surface_PT(id, &
           skip_partials, &
           lnT_surf, dlnT_dL, dlnT_dlnR, dlnT_dlnM, dlnT_dlnkap, &
           lnP_surf, dlnP_dL, dlnP_dlnR, dlnP_dlnM, dlnP_dlnkap, ierr)

use const_def, only: dp
use auto_diff
integer, intent(in) :: id
logical, intent(in) :: skip_partials
real(dp), intent(out) :: &
   lnT_surf, dlnT_dL, dlnT_dlnR, dlnT_dlnM, dlnT_dlnkap, &
   lnP_surf, dlnP_dL, dlnP_dlnR, dlnP_dlnM, dlnP_dlnkap
integer, intent(out) :: ierr

type(auto_diff_real_4var_order1) :: L, lnR, lnM, lnkap, M, R, kap, g, lnT, lnP
real(dp) :: g

! Set up auto_diff base variables.
! Note that the ordering of which variable goes in which d1val slot doesn't matter.
! We just need to use the same ordering when reading the results out.

L = s%L(1)
L%d1val1 = 1d0

lnR = log(s%R(1))
lnR%d1val2 = 1d0

lnM = log(s%M(1))
lnM%d1val3 = 1d0

lnkap = log(s%opacity(1))
lnkap%d1val4 = 1d0

! Calculate lnT and lnP

M = exp(lnM)
R = exp(lnR)
kap = exp(lnkap)

lnT = log(pow(L / (4d0 * pi * boltzsig * pow2(R))))

g = M / pow2(R)
lnP = log(g / kap)

! Read results.
! Note that the order matches the order we used to set up the base variables.
lnT_surf = lnT%val
dlnT_dL = lnT%d1val1
dlnT_dlnR = lnT%d1val2
dlnT_dlnM = lnT%d1val3
dlnT_dlnkap = lnT%d1val4

lnP_surf = lnP%val
dlnP_dL = lnP%d1val1
dlnP_dlnR = lnP%d1val2
dlnP_dlnM = lnP%d1val3
dlnP_dlnkap = lnP%d1val4


end subroutine other_surface_PT

Note that auto_diff does not support the operation real variable = auto_diff variable. If you attempt to set a real equal to an auto_diff variable you will get an error that this operation is not implemented. That’s why we read the value of lnT and lnP above using lnT%val and lnP%val. The reason this operation is not supported is that real variables cannot store derivative information, so information would be lost in that conversion. If this were done as an intermediate operation in a calculation (rather than during read out) the resulting derivatives that get read out at the end would be incorrect. It is easy to make this conversion by mistake, and there is no valid use case for this operation other than at read out, so we choose to throw a compiler error rather than allow it.

Another common failure mode is not including use auto_diff at the start of the hook. That will result in a compiler error saying that the relevant auto_diff types are not defined.

The above example with auto_diff looks more complicated than assigning derivatives by hand, but that’s only because we used a simple example. For instance the atmosphere model of Jermyn+2021 could be implemented as:

subroutine other_surface_PT(id, skip_partials, &
         lnT_surf, dlnT_dL, dlnT_dlnR, dlnT_dlnM, dlnT_dlnkap, &
         lnP_surf, dlnP_dL, dlnP_dlnR, dlnP_dlnM, dlnP_dlnkap, ierr)

    use auto_diff
    integer, intent(in) :: id
    logical, intent(in) :: skip_partials
      real(dp), intent(out) :: &
         lnT_surf, dlnT_dL, dlnT_dlnR, dlnT_dlnM, dlnT_dlnkap, &
         lnP_surf, dlnP_dL, dlnP_dlnR, dlnP_dlnM, dlnP_dlnkap
      integer, intent(out) :: ierr

   type (star_info), pointer :: s
   real(dp) :: minT, blend_minT, Ledd_rdp
   type(auto_diff_real_4var_order1) :: M, R, L, kappa
   type(auto_diff_real_4var_order1) :: R_ph, R_bondi, R_acc, vesc2, Mdot, loss, super_eddington_factor, L_incident
   type(auto_diff_real_4var_order1) :: H_disk, erf_arg, erf_f1, erf_f2, vortpar, R_AGN, R_Hill, tidepar
   type(auto_diff_real_4var_order1) :: L_stream, L_shock, Ledd, T_surf, P_surf, P_rad, P_acc, P_outflow, P_simple
   type(auto_diff_real_4var_order1) :: agn_Teff, simple_Teff, simple_lnT, simple_lnP, agn_lnT, agn_lnP
   type(auto_diff_real_4var_order1) :: lnT, lnP

   ! Basic setup
   ierr = 0
   call star_ptr(id, s, ierr)
   if (ierr /= 0) return

   ! Setup autodiff variables
   M = s%m(1)
   M%d1val1 = 1d0

   R = s%r(1)
   R%d1val2 = 1d0

   L = s%L(1)
   L%d1val3 = 1d0

   kappa = const_kappa
   kappa%d1val4 = 0d0
   !kappa = s%opacity(1)
   !kappa%d1val4 = 1d0

   ! Basic quantities
   R_bondi = 2d0*standard_cgrav * M / pow2(const_csb) ! Bondi radius
   R_acc = R_bondi
   Mdot = pi * pow2(R_bondi) * const_rho * const_csb ! Mdot gain

   R_ph = pow2(kappa * const_rho) * pow3(R_bondi) ! Photosphere radius
   R_ph = R_ph + R
   R_ph = min(R_ph, R_bondi)

   ! Adjust gain down due to radiation pressure
   !Ledd_rdp = eval_Ledd(s, ierr)
   !Ledd = (Ledd_rdp / M%val) * M ! Gives Ledd the correct derivatives, because Ledd_rdp scales like M.
   Ledd = 4.0d0*pi*clight*s%cgrav(1)*M/kappa

   if (s% x_logical_ctrl(7)) then
      Ledd = Ledd * max(1e-2, (1d0 - pow2(s%job%extras_rpar(i_mu_J) / s%job%extras_rpar(i_J_crit))))
   end if

   if (s% x_logical_ctrl(1)) then
      if (L < Ledd) then
         Mdot = Mdot * (1d0 - L/Ledd)**2d0 ! Suggested by Alexander Dittmann
      else
         Mdot = 0d0
      end if
   else
      Mdot = Mdot * (1d0 - tanh(abs(L / Ledd)))
   end if

   ! Adjust gain down due to gap opening in the disk
   if (s% x_logical_ctrl(2)) then
      Mdot = Mdot * exp(-pow(M / Msun / cutoff_mass, 2d0/3d0)) ! Form suggested by Doug Lin
   end if


   if (s% x_logical_ctrl(3)) then
      H_disk = pow(2d0, 0.5d0) * const_csb / Omega_AGN
      erf_arg = -R_bondi / H_disk
      erf_f1 = (-2d0/pow(pi, 0.5d0) )*pow( 1d0 - exp(-erf_arg*erf_arg), 0.5d0)
      erf_f2 = pow(pi, 0.5d0)*0.5d0 + (3.1d1/2d2)*exp(-erf_arg*erf_arg) - (3.41d2/8d3)*exp(-2d0*erf_arg*erf_arg)
      Mdot = 0.5d0 * Mdot * pow( pi, 0.5d0 ) * erf_f1 * erf_f2 / erf_arg
   end if

   if (s% x_logical_ctrl(4)) then
         R_AGN = pow( standard_cgrav * Mass_AGN * Msun / (Omega_AGN * Omega_AGN) , 1d0/3d0 )
         vortpar = (R_bondi * R_bondi * Omega_AGN)/(2d0 * const_csb * R_AGN) !avg vorticity parameter at R_bondi
         vortpar = (2d0 / (pi * vortpar))*asinh( pow(2d0*vortpar, 1d0/3d0))  !Krumholz, McKee, Klein 2005 eqn. A7
         Mdot = Mdot * pow(1d0 + pow(vortpar, -10d0), -0.1d0) !pick minimum factor
   end if

   if (s% x_logical_ctrl(5)) then
         R_Hill = pow(standard_cgrav * s%m(1) / (3d0 * Omega_AGN*Omega_AGN), 1d0/3d0)
         tidepar = pow(R_Hill / R_bondi, 2d0)
         Mdot = Mdot * min(1d0, tidepar) !pick minimum factor
         R_acc = min(R_Hill, R_Bondi)
         ! Use min(Bondi, Hill) for accretion rate. Similar to Rosenthal et al. 2020
   end if

   if (s% x_logical_ctrl(6)) then
         H_disk = pow(2d0, 0.5d0) * const_csb / Omega_AGN
         R_Hill = pow(standard_cgrav * s%m(1) / (3d0 * Omega_AGN*Omega_AGN), 1d0/3d0)
         R_acc = min(R_Hill, R_Bondi)
         tidepar = min(1d0, R_Hill/R_bondi)
         tidepar = min(tidepar, H_disk/R_bondi)
         tidepar = tidepar * min(1d0, R_Hill/R_bondi)
         tidepar = tidepar * min(1d0, exp(2d0*(1d0-pow(R_Hill/H_disk, 2d0))))
         !Dobbs-Dixon, Li, Lin 2007, approx eqn. 28
         Mdot = Mdot * tidepar
         ! Adjust accretion rate to take into account the tidal barrier
   end if

   ! Calculate super-eddington mass loss
   vesc2 = 2d0*standard_cgrav*M / R
   super_eddington_factor = 0.1d0
   loss = super_eddington_mdot_autodiff(Ledd,M,L,R,super_eddington_factor)

   ! Set a minimum on Teff.
   minT = const_Tb

   ! Various luminosities
   L_shock = Mdot * M * standard_cgrav / R
   L_stream = L + L_shock
   L_incident = 4 * pi * pow2(R_ph) * boltz_sigma * pow4(minT)
   L_stream = L_stream + L_incident
   L_stream = max(L_stream, L_incident) ! Avoid NaN's due to bad near-surface grad(T)

   ! Evaluate simple photosphere
   simple_Teff = pow(L_stream / (4d0 * pi * R**2 * boltz_sigma), 0.25d0)
   simple_lnP = log(standard_cgrav * M / (R**2 * kappa))
   P_simple = exp(simple_lnP)
   simple_lnT = log(simple_Teff) ! First cell is tau=2/3 in this simple solution.

   ! Evaluate complicated photosphere
   agn_Teff = pow(L_stream / (4 * pi * boltz_sigma * pow2(R_ph)), 0.25d0)
   T_surf = agn_Teff * pow(R_ph / R, 5d0 / 8d0)
   agn_lnT = log(T_surf)

   ! P_surf
   P_acc = const_rho * pow2(const_csb) * pow(R_acc / R, 2.5d0)
   P_rad = crad * pow4(T_surf) / 3d0
   P_outflow = - loss * pow(vesc2, 0.5d0) / (4 * pi * pow2(R))
   P_surf = P_acc + P_rad + P_outflow + P_simple
   agn_lnP = log(P_surf)

   ! blend gets set in extras_start_step
   lnT = agn_lnT
   lnP = agn_lnP

   ! Format for output
   lnT_surf = lnT%val
   dlnT_dlnM = lnT%d1val1 * M%val
   dlnT_dlnR = lnT%d1val2 * R%val
   dlnT_dL = lnT%d1val3
   dlnT_dlnkap = lnT%d1val4 * kappa%val

   lnP_surf = lnP%val
   dlnP_dlnM = lnP%d1val1 * M%val
   dlnP_dlnR = lnP%d1val2 * R%val
   dlnP_dL = lnP%d1val3
   dlnP_dlnkap = lnP%d1val4 * kappa%val

end subroutine other_surface_PT

Propagating the derivatives through this by hand would be tedious and error-prone. With auto_diff the focus stays on the physics rather than differentiation.